fixed mypy errors

Author: shiinamiyuki

luisa_lang/__init__.py

@@ -6,5 +6,4 @@
 
 from luisa_lang.lang import *
 from luisa_lang.lang_builtins import *
-
 bool = boolean

luisa_lang/ast_rewrite.py

@@ -474,7 +474,7 @@ def rewrite_function[F: Callable[..., Any]](f: F, decorator_name: str) -> F:
     tree, filename = retrieve_ast_and_filename(f)
     tree = FuncRewriter(decorator_name, filename).visit(tree)
     ast.fix_missing_locations(tree)
-    print(ast.unparse(tree))
+    # print(ast.unparse(tree))
     code = compile(tree, filename="<ast>", mode="exec")
     local_dict: dict[Any, Any] = {}
     exec(code, f.__globals__, local_dict)

luisa_lang/lang_builtins.py

@@ -1,6 +1,6 @@
 import types
 
-from _builtin_decor import func
+from luisa_lang._builtin_decor import func
 from luisa_lang.math_types import *
 from luisa_lang.core_types import Ref
 import luisa_lang.hir as hir

luisa_lang/lang_runtime.py

@@ -7,8 +7,8 @@
 
 from luisa_lang.utils import IdentityDict, check_type, is_generic_class
 import luisa_lang.hir as hir
-from hir import PyTreeStructure
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union, cast
+from luisa_lang.hir import PyTreeStructure
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union, cast
 
 
 class Scope:
@@ -174,11 +174,11 @@ def __init__(self, node: hir.Value, scope: Scope | None = None):
 
 
 class FlattenedTree:
-    metadata: Tuple[type, Tuple[Any], Any]  # (type, type_args, Any)
+    metadata: Tuple[Type[Any], Tuple[Any], Any]  # (type, type_args, Any)
     children: List["FlattenedTree"]
 
     def __init__(
-        self, metadata: Tuple[type, Tuple[Any], Any], children: List["FlattenedTree"]
+        self, metadata: Tuple[Type[Any], Tuple[Any], Any], children: List["FlattenedTree"]
     ):
         self.metadata = metadata
         self.children = children
@@ -342,13 +342,13 @@ def tree_flatten(obj: Any, allow_non_pytree_objects: bool) -> FlattenedTree:
     if isinstance(obj, PyTree):
         return obj._flatten()
     if allow_non_pytree_objects and not PyTreeRegistry.is_registered(type(obj)):
-        return FlattenedTree((type(obj), tuple(), obj), [])
+        return FlattenedTree((type(obj), cast(Tuple[Any, ...], tuple()), obj), [])
     flatten_func, _ = PyTreeRegistry.get(type(obj))
     return flatten_func(obj)
 
 
 def tree_unflatten(obj: FlattenedTree, allow_non_pytree_objects: bool) -> Any:
-    typ = obj.metadata[0]
+    typ: Type[Any] = obj.metadata[0]
     if issubclass(typ, JitVar):
         _type_args, v = obj.metadata[1:]
         assert isinstance(v, JitVar)
@@ -409,7 +409,7 @@ def is_registered(typ: type) -> bool:
     @staticmethod
     def __register_default_types() -> None:
         def flatten_primitive(obj: Any) -> FlattenedTree:
-            return FlattenedTree((type(obj), tuple(), obj), [])
+            return FlattenedTree((type(obj), cast(Tuple[Any, ...], tuple()), obj), [])
 
         def unflatten_primitive(tree: FlattenedTree) -> Any:
             assert len(tree.children) == 0
@@ -423,7 +423,7 @@ def unflatten_primitive(tree: FlattenedTree) -> Any:
 
         def flatten_list(obj: List[Any]) -> FlattenedTree:
             return FlattenedTree(
-                (list, tuple(), None), [tree_flatten(o, True) for o in obj]
+                (list, cast(Tuple[Any, ...], tuple()), None), [tree_flatten(o, True) for o in obj]
             )
 
         def unflatten_list(tree: FlattenedTree) -> List[Any]:
@@ -435,7 +435,7 @@ def unflatten_list(tree: FlattenedTree) -> List[Any]:
 
         def flatten_tuple(obj: Tuple[Any, ...]) -> FlattenedTree:
             return FlattenedTree(
-                (tuple, tuple(), None), [tree_flatten(o, True) for o in obj]
+                (tuple, cast(Tuple[Any, ...], tuple()), None), [tree_flatten(o, True) for o in obj]
             )
 
         def unflatten_tuple(tree: FlattenedTree) -> Tuple[Any, ...]:
@@ -447,14 +447,15 @@ def unflatten_tuple(tree: FlattenedTree) -> Tuple[Any, ...]:
 
         def flatten_dict(obj: Dict[Any, Any]) -> FlattenedTree:
             return FlattenedTree(
-                (dict, tuple(), (len(obj.keys()))),
+                (dict, cast(Tuple[Any, ...], tuple()), (len(obj.keys()))),
                 [tree_flatten(k, True) for k in obj.keys()]
                 + [tree_flatten(v, True) for v in obj.values()],
             )
 
         def unflatten_dict(tree: FlattenedTree) -> Dict[Any, Any]:
             assert tree.metadata[0] is dict
-            length = tree.metadata[1]
+            length = tree.metadata[2][0]
+            assert isinstance(length, int), "Invalid length for dict unflattening"
             assert len(tree.children) == length * 2
             keys = tree.children[:length]
             values = tree.children[length:]
@@ -561,7 +562,7 @@ class ControlFlowFrame:
     parent: Optional["ControlFlowFrame"]
     is_static: bool
 
-    def __init__(self, parent: Optional["ControlFlowFrame"]):
+    def __init__(self, *args, parent: Optional["ControlFlowFrame"]):
         self.parent = parent
         self.is_static = False
 
@@ -590,7 +591,7 @@ class IfFrame(ControlFlowFrame):
     false_bb: Optional[hir.BasicBlock]
 
     def __init__(self, cond: Any, parent: ControlFlowFrame):
-        super().__init__(parent)
+        super().__init__(parent=parent)
         self.cond = cond
         self.is_static = not isinstance(cond, JitVar)
         self.static_cond = bool(cond) if self.is_static else None
@@ -712,7 +713,7 @@ class TraceContext:
     top_level_func: Optional[hir.Function]
 
     def __init__(self, is_top_level):
-        self.cf_frame = ControlFlowFrame(None)
+        self.cf_frame = ControlFlowFrame(parent=None)
         self.is_top_level = is_top_level
         self.top_level_func = None