Merge pull request #6 from PyNumLab/codex/tighten-pyi-ast-parser

codex: tighten pyi AST parser helpers

Author: saidctb

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

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(
         """