diff --git a/semantics/pyi_parser.py b/semantics/pyi_parser.py index 2c9c90d30..4467242a9 100644 --- a/semantics/pyi_parser.py +++ b/semantics/pyi_parser.py @@ -67,32 +67,14 @@ def import_name(self, node: ast.Import) -> str: ) def class_def(self, node: ast.ClassDef, *, visibility: str) -> SemanticClass: - fields: list[SemanticArgument] = [] - methods: list[SemanticMethod] = [] - - for item in node.body: - if isinstance(item, ast.Pass): - continue - if isinstance(item, ast.AnnAssign): - fields.append(self.ann_assign(item, default_intent="in")) - continue - if isinstance(item, ast.FunctionDef): - decorators = self.decorators(item.decorator_list, context="class body") - methods.append( - self.method_def( - item, - visibility=decorators.visibility, - projection=decorators.projection, - ) - ) - continue - raise ValueError(f"Unsupported class body node: {_node_text(item)!r}") + body = _ClassBodyVisitor(self) + body.visit_body(node.body) return SemanticClass( name=node.name, native_name=node.name, - fields=fields, - methods=methods, + fields=body.fields, + methods=body.methods, base_classes=[ast.unparse(base) for base in node.bases], visibility=visibility, ) @@ -152,10 +134,10 @@ def ann_assign(self, node: ast.AnnAssign, *, default_intent: str) -> SemanticArg def decorators(self, nodes: list[ast.expr], *, context: str) -> _Decorators: parsed = _Decorators() for node in nodes: - if isinstance(node, ast.Name) and node.id == "private": + if self.matches_name(node, "private"): parsed.visibility = "private" continue - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "native_call": + if isinstance(node, ast.Call) and self.matches_name(node.func, "native_call"): parsed.has_native_call = True parsed.projection = self.native_call(node) continue @@ -174,12 +156,12 @@ def native_call(self, node: ast.Call) -> list[ProjectionMapping]: ] def native_projection_entry(self, node: ast.AST, native_position: int) -> ProjectionMapping: - if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): + if not isinstance(node, ast.Call): raise ValueError("native_call expects projection entry calls") if node.keywords: - raise ValueError(f"{node.func.id} expects positional arguments only") + raise ValueError(f"{self.required_name(node.func)} expects positional arguments only") - helper = node.func.id + helper = self.required_name(node.func) if helper == "Arg": if len(node.args) != 1: raise ValueError("Arg expects one positional index") @@ -239,27 +221,28 @@ def native_projection_entry(self, node: ast.AST, native_position: int) -> Projec raise ValueError(f"Unsupported native_call projection entry: {helper}") def native_value_ref(self, node: ast.AST) -> dict[str, int | str]: - if not isinstance(node, ast.Call) or not isinstance(node.func, ast.Name): + if not isinstance(node, ast.Call): raise ValueError("Expected Arg(...), Return(...), or Work(...) value reference") if node.keywords or len(node.args) != 1: - raise ValueError(f"{node.func.id} value reference expects one positional argument") - if node.func.id == "Arg": + raise ValueError(f"{self.required_name(node.func)} value reference expects one positional argument") + helper = self.required_name(node.func) + if helper == "Arg": return {"kind": "arg", "position": int(ast.literal_eval(node.args[0]))} - if node.func.id == "Return": + if helper == "Return": return {"kind": "return", "position": int(ast.literal_eval(node.args[0]))} - if node.func.id == "Work": + if helper == "Work": return {"kind": "work", "name": str(ast.literal_eval(node.args[0]))} raise ValueError("Expected Arg(...), Return(...), or Work(...) value reference") def visible_type(self, node: ast.expr) -> tuple[str, SemanticType, str | None]: - if self.subscript_name(node) == "private": + if self.is_subscript_of(node, "private"): semantic_type, original_name = self.semantic_type_annotation(self.subscript_slice(node)) return "private", semantic_type, original_name semantic_type, original_name = self.semantic_type_annotation(node) return "public", semantic_type, original_name def semantic_type_annotation(self, node: ast.expr) -> tuple[SemanticType, str | None]: - if self.subscript_name(node) != "Annotated": + if not self.is_subscript_of(node, "Annotated"): return self.semantic_type(node), None items = self.subscript_items(node) @@ -274,7 +257,7 @@ def semantic_type_annotation(self, node: ast.expr) -> tuple[SemanticType, str | return self.semantic_type(items[0]), original_name def semantic_type(self, node: ast.expr) -> SemanticType: - if self.subscript_name(node) == "Annotated": + if self.is_subscript_of(node, "Annotated"): semantic_type, _ = self.semantic_type_annotation(node) return semantic_type @@ -299,9 +282,9 @@ def semantic_type(self, node: ast.expr) -> SemanticType: def constraint(self, node: ast.expr) -> SemanticConstraint: if isinstance(node, ast.Name): return SemanticConstraint(node.id) - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name): + if isinstance(node, ast.Call): return SemanticConstraint( - name=node.func.id, + name=self.required_name(node.func), arguments=[ast.literal_eval(arg) for arg in node.args], ) raise ValueError(f"Unsupported semantic type constraint: {ast.unparse(node)!r}") @@ -338,7 +321,7 @@ def return_projection(self, node: ast.expr) -> tuple[SemanticType | None, list[S return return_type, returned_args def returned_argument(self, node: ast.expr) -> SemanticArgument | None: - if self.subscript_name(node) != "Returns": + if not self.is_subscript_of(node, "Returns"): return None items = self.subscript_items(node) if len(items) not in {2, 3}: @@ -355,7 +338,7 @@ def returned_argument(self, node: ast.expr) -> SemanticArgument | None: @staticmethod def name_metadata(node: ast.expr) -> str | None: - if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == "Name": + if isinstance(node, ast.Call) and _PyiAstParser.matches_name(node.func, "Name"): if len(node.args) != 1: raise ValueError(f"Name metadata expects one argument: {ast.unparse(node)!r}") return str(ast.literal_eval(node.args[0])) @@ -374,10 +357,31 @@ def default_marks_optional(node: ast.expr | None) -> bool: return isinstance(node, ast.Constant) and node.value in {Ellipsis, None} @staticmethod - def subscript_name(node: ast.AST) -> str: - if isinstance(node, ast.Subscript): - return ast.unparse(node.value) - return "" + def qualified_name(node: ast.AST) -> tuple[str, ...] | None: + if isinstance(node, ast.Name): + return (node.id,) + if isinstance(node, ast.Attribute): + parent = _PyiAstParser.qualified_name(node.value) + if parent is None: + return None + return (*parent, node.attr) + return None + + @staticmethod + def matches_name(node: ast.AST, name: str) -> bool: + qualified = _PyiAstParser.qualified_name(node) + return qualified is not None and qualified[-1] == name + + @staticmethod + def required_name(node: ast.AST) -> str: + qualified = _PyiAstParser.qualified_name(node) + if qualified is None: + raise ValueError(f"Expected named helper: {ast.unparse(node)!r}") + return qualified[-1] + + @staticmethod + def is_subscript_of(node: ast.AST, name: str) -> bool: + return isinstance(node, ast.Subscript) and _PyiAstParser.matches_name(node.value, name) @staticmethod def subscript_slice(node: ast.AST) -> ast.expr: @@ -527,11 +531,41 @@ def _apply_native_call_argument_names( mapping.result_position = return_positions.get(arg.name) def return_items(self, node: ast.expr) -> list[ast.expr]: - if self.subscript_name(node) == "tuple": + if self.is_subscript_of(node, "tuple") or self.is_subscript_of(node, "Tuple"): return self.subscript_items(node) return [node] +class _ClassBodyVisitor(ast.NodeVisitor): + def __init__(self, parser: _PyiAstParser): + self.parser = parser + self.fields: list[SemanticArgument] = [] + self.methods: list[SemanticMethod] = [] + + def visit_body(self, nodes: list[ast.stmt]) -> None: + for node in nodes: + self.visit(node) + + def visit_Pass(self, node: ast.Pass) -> None: + return None + + def visit_AnnAssign(self, node: ast.AnnAssign) -> None: + self.fields.append(self.parser.ann_assign(node, default_intent="in")) + + def visit_FunctionDef(self, node: ast.FunctionDef) -> None: + decorators = self.parser.decorators(node.decorator_list, context="class body") + self.methods.append( + self.parser.method_def( + node, + visibility=decorators.visibility, + projection=decorators.projection, + ) + ) + + def generic_visit(self, node: ast.AST) -> None: + raise ValueError(f"Unsupported class body node: {_node_text(node)!r}") + + class _ModuleVisitor(ast.NodeVisitor): def __init__(self, parser: _PyiAstParser): self.parser = parser diff --git a/tests/pyi/test_pyi_to_ir.py b/tests/pyi/test_pyi_to_ir.py index a57b76117..a9f1d7a66 100644 --- a/tests/pyi/test_pyi_to_ir.py +++ b/tests/pyi/test_pyi_to_ir.py @@ -141,6 +141,26 @@ def test_pyi_parser_ignores_unknown_annotation_metadata(): assert module.variables[1].name == "native_alias" +def test_parse_pyi_text_accepts_qualified_ast_wrapper_names(): + module = parse_pyi_text( + """ +import typing + +alias: typing.Annotated[Float64[typing.Shape("1:n")], typing.Name("native_alias")] + +def f() -> typing.Tuple[Float64, typing.Returns["y", Float64]]: ... +""", + module_name="edited", + ) + + assert module.variables[0].name == "native_alias" + assert module.variables[0].semantic_type.shape == ["1:n"] + assert module.functions[0].return_type is not None + assert module.functions[0].return_type.name == "Float64" + assert module.functions[0].arguments[0].name == "y" + assert module.functions[0].arguments[0].intent == "out" + + def test_parse_pyi_text_accepts_ast_only_projection_value_refs(): module = parse_pyi_text( """