feat: handle import statements

import statements now introduce variables with unknown types to avoid "Unknown variable" warnings

gen/python.py

@@ -24,6 +24,17 @@ class ParamSpec:
         return self.pos + self.mixed + self.kw
 
 
+@dataclass(frozen=True, kw_only=True)
+class ImportAlias:
+    location: Location
+    name: str
+    alias: Optional[str] = None
+
+    @property
+    def imported_name(self) -> str:
+        return self.alias if self.alias is not None else self.name
+
+
 ###<
 
 
@@ -99,6 +110,16 @@ class ForStmt:
     body: list[Stmt]
 
 
+class ImportStmt:
+    imports: list[ImportAlias]
+
+
+class FromImportStmt:
+    module: Optional[str]
+    imports: list[ImportAlias]
+    level: int
+
+
 class RawStmt:
     stmt: ast.stmt
 

midas/ast/python.py

@@ -25,6 +25,17 @@ class ParamSpec:
         return self.pos + self.mixed + self.kw
 
 
+@dataclass(frozen=True, kw_only=True)
+class ImportAlias:
+    location: Location
+    name: str
+    alias: Optional[str] = None
+
+    @property
+    def imported_name(self) -> str:
+        return self.alias if self.alias is not None else self.name
+
+
 ####################
 # Type annotations #
 ####################
@@ -123,6 +134,12 @@ class Stmt(ABC):
         @abstractmethod
         def visit_for_stmt(self, stmt: ForStmt) -> T: ...
 
+        @abstractmethod
+        def visit_import_stmt(self, stmt: ImportStmt) -> T: ...
+
+        @abstractmethod
+        def visit_from_import_stmt(self, stmt: FromImportStmt) -> T: ...
+
         @abstractmethod
         def visit_raw_stmt(self, stmt: RawStmt) -> T: ...
 
@@ -207,6 +224,24 @@ class ForStmt(Stmt):
         return visitor.visit_for_stmt(self)
 
 
+@dataclass(frozen=True)
+class ImportStmt(Stmt):
+    imports: list[ImportAlias]
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_import_stmt(self)
+
+
+@dataclass(frozen=True)
+class FromImportStmt(Stmt):
+    module: Optional[str]
+    imports: list[ImportAlias]
+    level: int
+
+    def accept(self, visitor: Stmt.Visitor[T]) -> T:
+        return visitor.visit_from_import_stmt(self)
+
+
 @dataclass(frozen=True)
 class RawStmt(Stmt):
     stmt: ast.stmt

midas/checker/python.py

@@ -603,6 +603,23 @@ class PythonTyper(
         if body_returned:
             raise ReturnException()
 
+    def visit_import_stmt(self, stmt: p.ImportStmt) -> None:
+        self._visit_imports(stmt.location, stmt.imports)
+
+    def visit_from_import_stmt(self, stmt: p.FromImportStmt) -> None:
+        self._visit_imports(stmt.location, stmt.imports)
+
+    def _visit_imports(self, location: Location, imports: list[p.ImportAlias]) -> None:
+        for import_ in imports:
+            self._assign_var(
+                location,
+                p.VariableExpr(
+                    name=import_.imported_name,
+                    location=import_.location,
+                ),
+                UnknownType(),
+            )
+
     def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
         pass
 

midas/checker/resolver.py

@@ -175,6 +175,18 @@ class Resolver(p.Stmt.Visitor[None], p.Expr.Visitor[None]):
         self.resolve(*stmt.body)
         self.end_scope()
 
+    def visit_import_stmt(self, stmt: p.ImportStmt) -> None:
+        self._resolve_imports(stmt.imports)
+
+    def visit_from_import_stmt(self, stmt: p.FromImportStmt) -> None:
+        self._resolve_imports(stmt.imports)
+
+    def _resolve_imports(self, imports: list[p.ImportAlias]) -> None:
+        for import_ in imports:
+            name: str = import_.imported_name
+            self.declare(name)
+            self.define(name)
+
     def visit_raw_stmt(self, stmt: p.RawStmt) -> None:
         pass
 

midas/generator/generator.py

@@ -307,6 +307,27 @@ class Generator(p.Stmt.Visitor[ast.stmt], p.Expr.Visitor[ast.expr]):
             orelse=[],
         )
 
+    def visit_import_stmt(self, stmt: p.ImportStmt) -> ast.stmt:
+        return ast.Import(
+            names=self._convert_imports(stmt.imports),
+        )
+
+    def visit_from_import_stmt(self, stmt: p.FromImportStmt) -> ast.stmt:
+        return ast.ImportFrom(
+            module=stmt.module,
+            names=self._convert_imports(stmt.imports),
+            level=stmt.level,
+        )
+
+    def _convert_imports(self, imports: list[p.ImportAlias]) -> list[ast.alias]:
+        return [
+            ast.alias(
+                name=import_.name,
+                asname=import_.alias,
+            )
+            for import_ in imports
+        ]
+
     def visit_raw_stmt(self, stmt: p.RawStmt) -> ast.stmt:
         return stmt.stmt
 

midas/parser/python.py

@@ -16,9 +16,12 @@ from midas.ast.python import (
     ForStmt,
     FrameColumn,
     FrameType,
+    FromImportStmt,
     Function,
     GetExpr,
     IfStmt,
+    ImportAlias,
+    ImportStmt,
     ListExpr,
     LiteralExpr,
     LogicalExpr,
@@ -105,10 +108,34 @@ class PythonParser:
             case ast.For(orelse=[]):
                 return self.parse_for(node)
 
+            case ast.Import(names=imports):
+                return ImportStmt(
+                    location=location,
+                    imports=self._parse_imports(imports),
+                )
+
+            case ast.ImportFrom(module=module, names=imports, level=level):
+                return FromImportStmt(
+                    location=location,
+                    module=module,
+                    imports=self._parse_imports(imports),
+                    level=level,
+                )
+
             case _:
                 print(f"Unsupported statement: {ast.unparse(node)}")
                 return RawStmt(location=location, stmt=node)
 
+    def _parse_imports(self, imports: list[ast.alias]) -> list[ImportAlias]:
+        return [
+            ImportAlias(
+                location=Location.from_ast(import_),
+                name=import_.name,
+                alias=import_.asname,
+            )
+            for import_ in imports
+        ]
+
     def parse_annotation_assign(self, node: ast.AnnAssign) -> list[Stmt]:
         statements: list[Stmt] = []
         loc: Location = Location.from_ast(node)

tests/__main__.py

@@ -31,7 +31,7 @@ def main():
         GeneratorTester,
     ]
 
-    success: bool = all(map(run_tests, testers))
+    success: bool = all(list(map(run_tests, testers)))  # list to avoid early stop
 
     if success:
         print(Ansi.FG(Ansi.BRIGHT_GREEN) + "All tests passed!" + Ansi.RESET)

tests/cases/python-parser/01_simple_types.py.ref.json

@@ -1,8 +1,16 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
+      "_type": "FromImportStmt",
-      "stmt": "from __future__ import annotations"
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "TypeAssign",

tests/cases/python-parser/02_custom_types.py.ref.json

@@ -1,8 +1,16 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
+      "_type": "FromImportStmt",
-      "stmt": "from __future__ import annotations"
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "TypeAssign",

tests/cases/python-parser/03_functions.py.ref.json

@@ -1,8 +1,16 @@
 {
   "stmts": [
     {
-      "_type": "RawStmt",
+      "_type": "FromImportStmt",
-      "stmt": "from __future__ import annotations"
+      "module": "__future__",
+      "imports": [
+        {
+          "_type": "ImportAlias",
+          "name": "annotations",
+          "alias": null
+        }
+      ],
+      "level": 0
     },
     {
       "_type": "Function",

tests/serializer/python.py

@@ -15,9 +15,12 @@ from midas.ast.python import (
     ForStmt,
     FrameColumn,
     FrameType,
+    FromImportStmt,
     Function,
     GetExpr,
     IfStmt,
+    ImportAlias,
+    ImportStmt,
     ListExpr,
     LiteralExpr,
     LogicalExpr,
@@ -194,6 +197,30 @@ class PythonAstJsonSerializer(
             "body": self._serialize_list(stmt.body),
         }
 
+    def visit_import_stmt(self, stmt: ImportStmt) -> dict:
+        return {
+            "_type": "ImportStmt",
+            "imports": self._serialize_imports(stmt.imports),
+        }
+
+    def visit_from_import_stmt(self, stmt: FromImportStmt) -> dict:
+        return {
+            "_type": "FromImportStmt",
+            "module": stmt.module,
+            "imports": self._serialize_imports(stmt.imports),
+            "level": stmt.level,
+        }
+
+    def _serialize_imports(self, imports: list[ImportAlias]) -> list:
+        return [
+            {
+                "_type": "ImportAlias",
+                "name": import_.name,
+                "alias": import_.alias,
+            }
+            for import_ in imports
+        ]
+
     def visit_raw_stmt(self, stmt: RawStmt) -> dict:
         return {
             "_type": "RawStmt",