simple type parsing works

luisa_lang/_utils.py

@@ -54,4 +54,6 @@ def retrieve_ast_and_filename(f: object) -> Tuple[ast.AST, str]:
     source_lines, lineno = sourceinspect.getsourcelines(f)
     src, indent = dedent_and_retrieve_indentation(source_lines)
     tree = increment_lineno_and_col_offset(ast.parse(src), lineno - 1, indent + 1)
+    for child in ast.walk(tree):
+        setattr(child, "source_file", source_file)
     return tree, source_file

luisa_lang/hir.py

@@ -201,7 +201,7 @@ def __eq__(self, value: object) -> bool:
             and value.element == self.element
             and value.count == self.count
         )
-    
+
     def __repr__(self) -> str:
         return f"ArrayType({self.element}, {self.count})"
 
@@ -344,8 +344,11 @@ def from_ast(ast: ast.AST) -> Optional["Span"]:
             return None
         if not hasattr(ast, "end_col_offset") or ast.end_col_offset is None:
             return None
+        file = None
+        if hasattr(ast, "source_file"):
+            file = getattr(ast, "source_file")
         return Span(
-            file=None,
+            file=file,
             start=(getattr(ast, "lineno", 0), getattr(ast, "col_offset", 0)),
             end=(getattr(ast, "end_lineno", 0), getattr(ast, "end_col_offset", 0)),
         )
@@ -554,7 +557,7 @@ def get() -> "GlobalContext":
 
     def __init__(self) -> None:
         assert _global_context is None, "GlobalContext should be a singleton"
-        self.types = {}
+        self.types = {type(None): UnitType()}
         self.functions = {}
 
 

luisa_lang/parse.py

@@ -270,7 +270,7 @@ def __init__(self, func: ast.FunctionDef, p_ctx: ParsingContext) -> None:
         self.arg_types = []
         self.return_type = None
         self._init_signature()
-        print(self.arg_types, "->", self.return_type)
+        print(self.arg_types, '->', self.return_type)
 
     def _init_signature(
         self,
@@ -288,6 +288,8 @@ def _init_signature(
                 raise RuntimeError("TODO: infer type")
             if (arg_ty := p_ctx.parse_type(arg.annotation)) is not None:
                 self.arg_types.append(arg_ty)
+            else:
+                report_error(arg.annotation, f"invalid type for argument")
         if func.returns is None:
             self.return_type = hir.UnitType()
         else:

scripts/gen_math_types.py

@@ -176,10 +176,7 @@ class {ty}:
                 name = names[bits]
                 gen_vector_type(f"{name}{size}", f"i{bits}", "int", size)
                 gen_vector_type(f"u{name}{size}", f"u{bits}", "int", size)
-    with open("luisa_lang/_math_type_exports.py", "w") as f:
-        __builtins__.print(
-            f'from luisa_lang.math_types import ({", ".join(exports)})', file=f
-        )
+        print("__all__ = " + str(exports))
 
 
 if __name__ == "__main__":