feat: handle import statements
import statements now introduce variables with unknown types to avoid "Unknown variable" warnings
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user