From 3de36e5bd54e27a118394f3f5ef9f764ff8197b4 Mon Sep 17 00:00:00 2001 From: LordBaryhobal Date: Mon, 6 Jul 2026 15:21:56 +0200 Subject: [PATCH] feat: handle import statements import statements now introduce variables with unknown types to avoid "Unknown variable" warnings --- gen/python.py | 21 +++++++++++ midas/ast/python.py | 35 +++++++++++++++++++ midas/checker/python.py | 17 +++++++++ midas/checker/resolver.py | 12 +++++++ midas/generator/generator.py | 21 +++++++++++ midas/parser/python.py | 27 ++++++++++++++ tests/__main__.py | 2 +- .../python-parser/01_simple_types.py.ref.json | 12 +++++-- .../python-parser/02_custom_types.py.ref.json | 12 +++++-- .../python-parser/03_functions.py.ref.json | 12 +++++-- tests/serializer/python.py | 27 ++++++++++++++ 11 files changed, 191 insertions(+), 7 deletions(-) diff --git a/gen/python.py b/gen/python.py index 60a92f7..6a3b1eb 100644 --- a/gen/python.py +++ b/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 diff --git a/midas/ast/python.py b/midas/ast/python.py index 65e896c..9b48552 100644 --- a/midas/ast/python.py +++ b/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 diff --git a/midas/checker/python.py b/midas/checker/python.py index 674727d..8ed74ad 100644 --- a/midas/checker/python.py +++ b/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 diff --git a/midas/checker/resolver.py b/midas/checker/resolver.py index d7af800..deef1cb 100644 --- a/midas/checker/resolver.py +++ b/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 diff --git a/midas/generator/generator.py b/midas/generator/generator.py index 07799b8..b31164a 100644 --- a/midas/generator/generator.py +++ b/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 diff --git a/midas/parser/python.py b/midas/parser/python.py index 966ab16..269000e 100644 --- a/midas/parser/python.py +++ b/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) diff --git a/tests/__main__.py b/tests/__main__.py index 0280e3b..f80d426 100644 --- a/tests/__main__.py +++ b/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) diff --git a/tests/cases/python-parser/01_simple_types.py.ref.json b/tests/cases/python-parser/01_simple_types.py.ref.json index 50730f3..1dbaa9e 100644 --- a/tests/cases/python-parser/01_simple_types.py.ref.json +++ b/tests/cases/python-parser/01_simple_types.py.ref.json @@ -1,8 +1,16 @@ { "stmts": [ { - "_type": "RawStmt", - "stmt": "from __future__ import annotations" + "_type": "FromImportStmt", + "module": "__future__", + "imports": [ + { + "_type": "ImportAlias", + "name": "annotations", + "alias": null + } + ], + "level": 0 }, { "_type": "TypeAssign", diff --git a/tests/cases/python-parser/02_custom_types.py.ref.json b/tests/cases/python-parser/02_custom_types.py.ref.json index fdb89d3..e59716d 100644 --- a/tests/cases/python-parser/02_custom_types.py.ref.json +++ b/tests/cases/python-parser/02_custom_types.py.ref.json @@ -1,8 +1,16 @@ { "stmts": [ { - "_type": "RawStmt", - "stmt": "from __future__ import annotations" + "_type": "FromImportStmt", + "module": "__future__", + "imports": [ + { + "_type": "ImportAlias", + "name": "annotations", + "alias": null + } + ], + "level": 0 }, { "_type": "TypeAssign", diff --git a/tests/cases/python-parser/03_functions.py.ref.json b/tests/cases/python-parser/03_functions.py.ref.json index 64e8c00..a4aafdf 100644 --- a/tests/cases/python-parser/03_functions.py.ref.json +++ b/tests/cases/python-parser/03_functions.py.ref.json @@ -1,8 +1,16 @@ { "stmts": [ { - "_type": "RawStmt", - "stmt": "from __future__ import annotations" + "_type": "FromImportStmt", + "module": "__future__", + "imports": [ + { + "_type": "ImportAlias", + "name": "annotations", + "alias": null + } + ], + "level": 0 }, { "_type": "Function", diff --git a/tests/serializer/python.py b/tests/serializer/python.py index afb5521..8bd73a7 100644 --- a/tests/serializer/python.py +++ b/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",