feat: handle import statements

import statements now introduce variables with unknown types to avoid "Unknown variable" warnings
This commit is contained in:
2026-07-06 15:21:56 +02:00
parent 7075d011b8
commit 3de36e5bd5
11 changed files with 191 additions and 7 deletions

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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)

View File

@@ -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)

View File

@@ -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",

View File

@@ -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",

View File

@@ -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",

View File

@@ -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",