[bug] Allow StructType as type hint to ti.func

taichi-dev · Dec 22, 2022 · 8538a61 · 8538a61
1 parent be15358
commit 8538a61
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/python/taichi/lang/kernel_impl.py b/python/taichi/lang/kernel_impl.py
@@ -334,6 +334,8 @@ def extract_arguments(self):
                     pass
                 elif isinstance(annotation, MatrixType):
                     pass
+                elif isinstance(annotation, StructType):
+                    pass
                 elif id(annotation) in primitive_types.type_ids:
                     pass
                 elif isinstance(annotation, template):

diff --git a/tests/python/test_function.py b/tests/python/test_function.py
@@ -431,6 +431,24 @@ def test_error():
         test_error()
 
 
+@test_utils.test(debug=True)
+def test_func_struct_arg():
+    @ti.dataclass
+    class C:
+        i: int
+
+    @ti.func
+    def f(c: C):
+        return c.i
+
+    @ti.kernel
+    def k():
+        c = C(i=2)
+        assert f(c) == 2
+
+    k()
+
+
 @test_utils.test(arch=[ti.cpu, ti.cuda])
 def test_real_func_matrix_arg():
     @ti.experimental.real_func