frontend: (pyast) support classmethods (#5595)

Author: Time0o

tests/filecheck/frontend/programs/registered_classmethods.py

@@ -0,0 +1,107 @@
+# RUN: python %s | filecheck %s
+
+from ctypes import c_int32
+
+from xdsl.dialects import arith, builtin
+from xdsl.frontend.pyast.context import PyASTContext
+from xdsl.frontend.pyast.utils.exceptions import CodeGenerationException
+
+
+class Adder:
+    @classmethod
+    def add_i32(cls, operand1: c_int32, operand2: c_int32) -> c_int32: ...
+
+
+ctx = PyASTContext()
+ctx.register_type(c_int32, builtin.i32)
+ctx.register_function(Adder.add_i32, arith.AddiOp)
+
+
+@ctx.parse_program
+def test_add(x: c_int32, y: c_int32) -> c_int32:
+    return Adder.add_i32(x, operand2=y)
+
+
+print(test_add.module)
+# CHECK-NEXT: builtin.module {
+# CHECK-NEXT:   func.func @test_add(%x : i32, %y : i32) -> i32 {
+# CHECK-NEXT:     %0 = arith.addi %x, %y : i32
+# CHECK-NEXT:     func.return %0 : i32
+# CHECK-NEXT:   }
+# CHECK-NEXT: }
+
+
+# CHECK-NEXT: Classmethod arguments must be declared variables.
+@ctx.parse_program
+def test_args():
+    return Adder.add_i32(1, 2)  # pyright: ignore[reportArgumentType]
+
+
+try:
+    test_args.module
+except CodeGenerationException as e:
+    print(e.msg)
+
+
+# ================================================= #
+# Disable the desymref pass for the remaining tests #
+# ================================================= #
+ctx.post_transforms = []
+
+
+# CHECK-NEXT: Classmethod arguments must be declared variables.
+@ctx.parse_program
+def test_more_args():
+    return Adder.add_i32(operand1=1, operand2=2)  # pyright: ignore[reportArgumentType]
+
+
+try:
+    test_more_args.module
+except CodeGenerationException as e:
+    print(e.msg)
+
+
+class Class:
+    @classmethod
+    def method(cls):
+        pass
+
+
+# CHECK-NEXT: Classmethod 'Class.method' is not registered.
+@ctx.parse_program
+def test_unregistered():
+    return Class.method()  # noqa: F821
+
+
+try:
+    test_unregistered.module
+except CodeGenerationException as e:
+    print(e.msg)
+
+
+# CHECK-NEXT: Method 'method' is not defined on class 'Class'.
+@ctx.parse_program
+def test_missing_method():
+    return Class.method()  # noqa: F821
+
+
+del Class.method
+
+try:
+    test_missing_method.module
+except CodeGenerationException as e:
+    print(e.msg)
+
+
+# CHECK-NEXT: Class 'Class' is not defined in scope.
+@ctx.parse_program
+def test_missing_class():
+    return Class.method()  # noqa: F821
+
+
+del Class
+
+try:
+    test_missing_class.module
+except CodeGenerationException as e:
+    print(e.msg)

tests/filecheck/frontend/pypdl/pdl.py

@@ -27,23 +27,16 @@
 
 from xdsl.dialects import arith, builtin, pdl
 from xdsl.frontend import pypdl
-from xdsl.ir import Operation
 from xdsl.rewriter import Rewriter
 
-
-def erase_op(operation: Operation) -> None:
-    """Shim to avoid `Expr` AST node required for methods."""
-    return Rewriter.erase_op(operation)
-
-
 ctx = pypdl.PyPDLContext()
 ctx.register_type(arith.ConstantOp, pdl.OperationType())
-ctx.register_function(erase_op, pdl.EraseOp)
+ctx.register_function(Rewriter.erase_op, pdl.EraseOp)
 
 
 @ctx.parse_program
 def constant_replace(matched_operation: arith.ConstantOp):
-    erase_op(matched_operation)
+    Rewriter.erase_op(matched_operation)
 
 
 # Check that the DSL correctly rewrites on the xDSL data structures

xdsl/frontend/pyast/code_generation.py

@@ -1,6 +1,7 @@
 import ast
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from typing import cast
+from typing import Any, cast
 
 import xdsl.dialects.builtin as builtin
 import xdsl.dialects.cf as cf
@@ -193,26 +194,28 @@ def visit_BinOp(self, node: ast.BinOp) -> None:
         )
 
     def visit_Call(self, node: ast.Call) -> None:
-        # Resolve function
-        assert isinstance(node.func, ast.Name)
-        func_name = node.func.id
-        source_func = self.type_converter.globals.get(func_name, None)
-        if source_func is None:
-            raise CodeGenerationException(
-                self.file,
-                node.lineno,
-                node.col_offset,
-                f"Function '{func_name}' is not defined in scope.",
-            )
-        ir_op = self.type_converter.function_registry.get_operation_constructor(
-            source_func
-        )
+        match node.func:
+            case ast.Name():
+                source_kind = "function"
+                source, source_name = self._call_source_function(node)
+            case ast.Attribute():
+                source_kind = "classmethod"
+                source, source_name = self._call_source_classmethod(node)
+            case _:
+                raise CodeGenerationException(
+                    self.file,
+                    node.lineno,
+                    node.col_offset,
+                    "Unsupported call expression.",
+                )
+
+        ir_op = self.type_converter.function_registry.get_operation_constructor(source)
         if ir_op is None:
             raise CodeGenerationException(
                 self.file,
                 node.lineno,
                 node.col_offset,
-                f"Function '{func_name}' is not registered.",
+                f"{source_kind.capitalize()} '{source_name}' is not registered.",
             )
 
         # Resolve arguments
@@ -224,7 +227,7 @@ def visit_Call(self, node: ast.Call) -> None:
                     self.file,
                     node.lineno,
                     node.col_offset,
-                    "Function arguments must be declared variables.",
+                    f"{source_kind.capitalize()} arguments must be declared variables.",
                 )
             args.append(arg_op := symref.FetchOp(arg.id, self.symbol_table[arg.id]))
             self.inserter.insert_op(arg_op)
@@ -240,7 +243,7 @@ def visit_Call(self, node: ast.Call) -> None:
                     self.file,
                     node.lineno,
                     node.col_offset,
-                    "Function arguments must be declared variables.",
+                    f"{source_kind.capitalize()} arguments must be declared variables.",
                 )
             assert keyword.arg is not None
             kwargs[keyword.arg] = symref.FetchOp(
@@ -250,6 +253,50 @@ def visit_Call(self, node: ast.Call) -> None:
 
         self.inserter.insert_op(ir_op(*args, **kwargs))
 
+    # Get called function for a call expression.
+    def _call_source_function(self, node: ast.Call) -> tuple[Callable[..., Any], str]:
+        assert isinstance(node.func, ast.Name)
+
+        func_name = node.func.id
+        func = self.type_converter.globals.get(func_name, None)
+        if func is None:
+            raise CodeGenerationException(
+                self.file,
+                node.lineno,
+                node.col_offset,
+                f"Function '{func_name}' is not defined in scope.",
+            )
+        return func, func_name
+
+    # Get called classmethod for a call expression.
+    def _call_source_classmethod(
+        self, node: ast.Call
+    ) -> tuple[Callable[..., Any], str]:
+        assert isinstance(node.func, ast.Attribute)
+        assert isinstance(node.func.value, ast.Name)
+
+        class_name = node.func.value.id
+        method_name = node.func.attr
+        classmethod_name = f"{class_name}.{method_name}"
+
+        source_class = self.type_converter.globals.get(class_name, None)
+        if source_class is None:
+            raise CodeGenerationException(
+                self.file,
+                node.lineno,
+                node.col_offset,
+                f"Class '{class_name}' is not defined in scope.",
+            )
+        classmethod_ = getattr(source_class, method_name, None)
+        if classmethod_ is None:
+            raise CodeGenerationException(
+                self.file,
+                node.lineno,
+                node.col_offset,
+                f"Method '{method_name}' is not defined on class '{class_name}'.",
+            )
+        return classmethod_, classmethod_name
+
     def visit_Compare(self, node: ast.Compare) -> None:
         # Allow a single comparison only.
         if len(node.comparators) != 1 or len(node.ops) != 1: