`ImportStmt`/`ImportFromStmt` - Represents an import statement.

xmnlab commented 1 month ago

from gpt:

note: expression classes should inherit from Expr and statement classes should inherit from StatementType. we probably should rename StatementType to Statement in the future

Overview

ImportStmt: Represents an import statement like import module_name.
ImportFromStmt: Represents an import-from statement like from module_name import name [as alias].
AliasExpr: Represents an alias in an import statement, handling the as alias part.

We will design these classes to integrate seamlessly with your existing AST framework, following the structure and conventions you've established.

Step-by-Step Guide

1. Understanding the Existing Structure

Before designing the new classes, let's briefly review the relevant parts of your AST framework:

AST: Base class for all AST nodes.
ASTKind: Enum class for different kinds of AST nodes.
ASTNodes: Class for nodes that contain a list of AST nodes (e.g., modules, blocks).
Expr: Represents expressions in the AST.
StatementType: Base class for statements.

Given that import statements are statements in the language, they should inherit from StatementType or AST.

2. Updating `ASTKind` Enum

We need to add new kinds to the ASTKind enum to represent the import statements and aliases.

Code for Updating `ASTKind`

# In src/astx/base.py, within the ASTKind enum, add:

# import statements
ImportStmtKind = -700
ImportFromStmtKind = -701
AliasExprKind = -702

Explanation

ImportStmtKind: Kind for ImportStmt.
ImportFromStmtKind: Kind for ImportFromStmt.
AliasExprKind: Kind for AliasExpr.

3. Defining the `AliasExpr` Class

First, let's define the AliasExpr class to represent aliases in import statements.

Code for `AliasExpr` Class

# In src/astx/datatypes.py or a new module like src/astx/statements.py

from astx.base import (
    AST,
    ASTKind,
    SourceLocation,
    NO_SOURCE_LOCATION,
    ASTNodes,
)
from astx.types import ReprStruct
from typing import Optional
from public import public

@public
class AliasExpr(AST):
    """Represents an alias in an import statement."""

    name: str
    asname: Optional[str]

    def __init__(
        self,
        name: str,
        asname: Optional[str] = None,
        loc: SourceLocation = NO_SOURCE_LOCATION,
        parent: Optional[ASTNodes] = None,
    ) -> None:
        super().__init__(loc=loc, parent=parent)
        self.name = name
        self.asname = asname
        self.kind = ASTKind.AliasExprKind

    def __str__(self) -> str:
        """Return a string representation of the alias."""
        if self.asname:
            return f"{self.name} as {self.asname}"
        else:
            return self.name

    def get_struct(self, simplified: bool = False) -> ReprStruct:
        """Return the AST structure of the alias."""
        key = "Alias"
        value = {
            "name": self.name,
            "asname": self.asname,
        }
        return self._prepare_struct(key, value, simplified)

Explanation

Attributes:
- name: The original name being imported.
- asname: The alias name, if any.
Methods:
- __str__: Provides a readable string representation.
- get_struct: Returns the AST structure for serialization and visualization.
Inheritance: Inherits from AST since it's a fundamental node.

4. Defining the `ImportStmt` Class

Now, let's define the ImportStmt class to represent statements like import module_name.

Code for `ImportStmt` Class

@public
class ImportStmt(AST):
    """Represents an import statement."""

    names: list[AliasExpr]

    def __init__(
        self,
        names: list[AliasExpr],
        loc: SourceLocation = NO_SOURCE_LOCATION,
        parent: Optional[ASTNodes] = None,
    ) -> None:
        super().__init__(loc=loc, parent=parent)
        self.names = names
        self.kind = ASTKind.ImportStmtKind

    def __str__(self) -> str:
        """Return a string representation of the import statement."""
        names_str = ", ".join(str(name) for name in self.names)
        return f"import {names_str}"

    def get_struct(self, simplified: bool = False) -> ReprStruct:
        """Return the AST structure of the import statement."""
        key = "Import"
        value = [name.get_struct(simplified) for name in self.names]
        return self._prepare_struct(key, value, simplified)

Explanation

Attributes:
- names: A list of AliasExpr instances representing the modules or names being imported.
Methods:
- __str__: Provides a readable string representation.
- get_struct: Returns the AST structure.
Inheritance: Inherits from AST.

5. Defining the `ImportFromStmt` Class

Next, define the ImportFromStmt class to represent statements like from module_name import name [as alias].

Code for `ImportFromStmt` Class

@public
class ImportFromStmt(AST):
    """Represents an import-from statement."""

    module: Optional[str]
    names: list[AliasExpr]
    level: int  # Represents the level of relative import (number of dots)

    def __init__(
        self,
        module: Optional[str],
        names: list[AliasExpr],
        level: int = 0,
        loc: SourceLocation = NO_SOURCE_LOCATION,
        parent: Optional[ASTNodes] = None,
    ) -> None:
        super().__init__(loc=loc, parent=parent)
        self.module = module
        self.names = names
        self.level = level
        self.kind = ASTKind.ImportFromStmtKind

    def __str__(self) -> str:
        """Return a string representation of the import-from statement."""
        level_dots = "." * self.level
        module_str = f"{level_dots}{self.module}" if self.module else level_dots
        names_str = ", ".join(str(name) for name in self.names)
        return f"from {module_str} import {names_str}"

    def get_struct(self, simplified: bool = False) -> ReprStruct:
        """Return the AST structure of the import-from statement."""
        key = "ImportFrom"
        value = {
            "module": self.module,
            "level": self.level,
            "names": [name.get_struct(simplified) for name in self.names],
        }
        return self._prepare_struct(key, value, simplified)

Explanation

Attributes:
- module: The module name from which to import; can be None for relative imports like from . import name.
- names: A list of AliasExpr instances representing the names being imported.
- level: The level of relative import (e.g., from ..module import name would have level=2).
Methods:
- __str__: Provides a readable string representation, handling relative imports.
- get_struct: Returns the AST structure.
Inheritance: Inherits from AST.

6. Testing the Implementation

Example Usage

# Example usage

# Create aliases
alias1 = AliasExpr(name="math")
alias2 = AliasExpr(name="os", asname="operating_system")

# Create an import statement
import_stmt = ImportStmt(names=[alias1, alias2])

# Print the string representation
print(import_stmt)

# Print the AST structure
print(import_stmt.to_yaml())

# Create an import-from statement
alias3 = AliasExpr(name="path", asname="p")
import_from_stmt = ImportFromStmt(module="os", names=[alias3], level=0)

# Print the string representation
print(import_from_stmt)

# Print the AST structure
print(import_from_stmt.to_yaml())

Expected Output

import math, os as operating_system

Import:
- Alias:
    name: math
    asname: null
- Alias:
    name: os
    asname: operating_system

from os import path as p

ImportFrom:
  module: os
  level: 0
  names:
  - Alias:
      name: path
      asname: p

7. Integration with `ASTNodes`

If you have a module or block of code represented by an ASTNodes instance, you can include import statements in it.

Example

module_ast = ASTNodes(name="module")
module_ast.append(import_stmt)
module_ast.append(import_from_stmt)

# Print the module's AST structure
print(module_ast.to_yaml())

Expected Output

entry:
- Import:
  - Alias:
      name: math
      asname: null
  - Alias:
      name: os
      asname: operating_system
- ImportFrom:
    module: os
    level: 0
    names:
    - Alias:
        name: path
        asname: p

8. Handling Edge Cases

Relative Imports

Ensure that relative imports without a module name are handled correctly (e.g., from . import name).

Wildcard Imports

Handle the case where names contains a wildcard (e.g., from module import *).

Code Update for `ImportFromStmt`

# Modify the __init__ method to accept a wildcard

def __init__(
    self,
    module: Optional[str],
    names: Optional[list[AliasExpr]] = None,
    level: int = 0,
    loc: SourceLocation = NO_SOURCE_LOCATION,
    parent: Optional[ASTNodes] = None,
) -> None:
    super().__init__(loc=loc, parent=parent)
    self.module = module
    self.names = names or []
    self.level = level
    self.kind = ASTKind.ImportFromStmtKind

# Modify __str__ method

def __str__(self) -> str:
    level_dots = "." * self.level
    module_str = f"{level_dots}{self.module}" if self.module else level_dots
    if self.names:
        names_str = ", ".join(str(name) for name in self.names)
    else:
        names_str = "*"
    return f"from {module_str} import {names_str}"

# Modify get_struct method

def get_struct(self, simplified: bool = False) -> ReprStruct:
    value = {
        "module": self.module,
        "level": self.level,
        "names": (
            [name.get_struct(simplified) for name in self.names]
            if self.names
            else ["*"]
        ),
    }
    return self._prepare_struct("ImportFrom", value, simplified)

9. Full Code Snippets

Updated `astx/base.py`

# ... existing imports ...

@public
class ASTKind(Enum):
    """The expression kind class used for downcasting."""

    # ... existing kinds ...

    # Import statements
    ImportStmtKind = -700
    ImportFromStmtKind = -701
    AliasExprKind = -702

    # ... rest of the code ...

New Module `astx/statements.py`

Since import statements are statements, it might be logical to place them in a new module, statements.py.

# src/astx/statements.py

from astx.base import (
    AST,
    ASTKind,
    SourceLocation,
    NO_SOURCE_LOCATION,
    ASTNodes,
)
from astx.types import ReprStruct
from typing import Optional, List
from public import public

@public
class AliasExpr(AST):
    # ... as defined above ...

@public
class ImportStmt(AST):
    # ... as defined above ...

@public
class ImportFromStmt(AST):
    # ... as defined above ...

10. Integration with the Rest of the Framework

Ensure that your module imports and exports are updated to include the new classes.

In `astx/init.py`

from astx.statements import AliasExpr, ImportStmt, ImportFromStmt

__all__ = [
    # ... existing exports ...
    "AliasExpr",
    "ImportStmt",
    "ImportFromStmt",
]

Conclusion

By defining the AliasExpr, ImportStmt, and ImportFromStmt classes, you've extended your AST framework to handle import statements, including aliases. These classes integrate with your existing structure and support AST representation and visualization.

Next Steps

Testing: Write unit tests to ensure that the classes handle various import scenarios correctly, including edge cases like relative imports and wildcard imports.
Error Handling: Implement validation in the constructors to handle invalid inputs (e.g., empty module names when not allowed).
Documentation: Update your module's documentation to include these new classes and explain their usage.

Additional Considerations

1. Handling Future Imports

If your language supports future imports (e.g., from __future__ import division in Python), you may want to handle them explicitly.

Code Suggestion

You can add an attribute or subclass to handle future imports.

@public
class ImportFutureStmt(ImportFromStmt):
    """Represents a future import statement."""

    def __init__(
        self,
        names: list[AliasExpr],
        loc: SourceLocation = NO_SOURCE_LOCATION,
        parent: Optional[ASTNodes] = None,
    ) -> None:
        super().__init__(module="__future__", names=names, loc=loc, parent=parent)
        self.kind = ASTKind.ImportFromStmtKind  # Use the same kind or create a new one

2. Integration with Symbol Table

If your AST framework interacts with a symbol table or namespace management, you might need to update those components to handle import statements and aliases.

3. Extending `AliasExpr`

If your language allows importing multiple levels (e.g., import package.module), you may need to represent the full dotted name.

Code Update

@public
class AliasExpr(AST):
    name: str  # Change to support dotted names
    asname: Optional[str]

    def __init__(
        self,
        name: str,
        asname: Optional[str] = None,
        loc: SourceLocation = NO_SOURCE_LOCATION,
        parent: Optional[ASTNodes] = None,
    ) -> None:
        super().__init__(loc=loc, parent=parent)
        self.name = name  # Can be 'package.module'
        self.asname = asname
        self.kind = ASTKind.AliasExprKind

**4. Representing `Import* `in`ImportStmt`**

If your language allows import * outside of from statements, you might need to handle that in ImportStmt.

xmnlab commented 1 month ago

Examples of ImportFromStmt in Python:

>>> ast.dump(ast.parse("from abc import ABC"))
"Module(body=[ImportFrom(module='abc', names=[alias(name='ABC')], level=0)], type_ignores=[])"
>>> ast.dump(ast.parse("from abc import ABC as MyABC"))
"Module(body=[ImportFrom(module='abc', names=[alias(name='ABC', asname='MyABC')], level=0)], type_ignores=[])"
>>> ast.dump(ast.parse("from .. import annotations"))
"Module(body=[ImportFrom(names=[alias(name='annotations')], level=2)], type_ignores=[])"
>>> ast.dump(ast.parse("from . import annotations"))
"Module(body=[ImportFrom(names=[alias(name='annotations')], level=1)], type_ignores=[])"
>>> ast.dump(ast.parse("from .mypkg import class1"))
"Module(body=[ImportFrom(module='mypkg', names=[alias(name='class1')], level=1)], type_ignores=[])"

Example of ImportStmt in Python:

>>> ast.dump(ast.parse("import ABC as MyABC"))
"Module(body=[Import(names=[alias(name='ABC', asname='MyABC')])], type_ignores=[])"
>>> ast.dump(ast.parse("import matplotlib.pyplot as plt"))
"Module(body=[Import(names=[alias(name='matplotlib.pyplot', asname='plt')])], type_ignores=[])"

So probably we should add a class for Alias

xmnlab commented 1 month ago

solved by @apkrelling #118 thanks!

arxlang / astx