Fix type_info null errors (#294)

Author: oulgen

helion/_compiler/device_ir.py

@@ -683,9 +683,9 @@ def visit_Name(self, node: ast.Name) -> object:
             return self.scope[node.id]
         assert isinstance(node, ExtendedAST)
         type_info = node._type_info
-        assert type_info.origin.is_host()  # pyright: ignore[reportOptionalMemberAccess]
+        assert type_info is not None and type_info.origin.is_host()
         try:
-            return type_info.proxy()  # pyright: ignore[reportOptionalMemberAccess]
+            return type_info.proxy()
         except NotImplementedError:
             raise exc.CantReadOnDevice(type_info) from None
 
@@ -757,6 +757,7 @@ def visit_Assign(self, node: ast.Assign) -> None:
         ):
             raise exc.NonTensorSubscriptAssign(lhs_type, rhs_type)
         assert isinstance(target.value, ExtendedAST)
+        assert target.value._type_info is not None
         target_origin = target.value._type_info.origin  # pyright: ignore[reportOptionalMemberAccess]
         if not target_origin.is_host():
             # Get the variable name for the error message
@@ -781,7 +782,8 @@ def _assign_subscript(self, target: ast.Subscript, val: object) -> None:
             raise exc.NonTensorSubscriptAssign(lhs_type, type(val))
 
         assert isinstance(target.value, ExtendedAST)
-        target_origin = target.value._type_info.origin  # pyright: ignore[reportOptionalMemberAccess]
+        assert target.value._type_info is not None
+        target_origin = target.value._type_info.origin
         assert target_origin.is_host()
 
         return hl.store(
@@ -811,7 +813,7 @@ def visit_Subscript(self, node: ast.Subscript) -> object:
         value = node.value
         assert isinstance(value, ExtendedAST)
         type_info = value._type_info
-        if type_info.origin.is_host():  # pyright: ignore[reportOptionalMemberAccess]
+        if type_info is not None and type_info.origin.is_host():
             return hl.load(self.visit(value), self._subscript_slice_proxy(node.slice))  # pyright: ignore[reportArgumentType]
         return hl.subscript(self.visit(value), self._subscript_slice_proxy(node.slice))  # pyright: ignore[reportArgumentType]
 

helion/_compiler/generate_ast.py

@@ -191,14 +191,17 @@ def visit_For(self, node: ast.For) -> ast.AST | None:
                     for arg_node in iter_node.args:
                         assert not isinstance(arg_node, ast.Starred)
                         assert isinstance(arg_node, ExtendedAST)
-                        args.append(arg_node._type_info.proxy())  # pyright: ignore[reportOptionalMemberAccess]
+                        assert arg_node._type_info is not None
+                        args.append(arg_node._type_info.proxy())
                     for kwarg_node in iter_node.keywords:
                         assert kwarg_node.arg is not None
                         assert isinstance(kwarg_node.value, ExtendedAST)
-                        kwargs[kwarg_node.arg] = kwarg_node.value._type_info.proxy()  # pyright: ignore[reportOptionalMemberAccess]
+                        assert kwarg_node.value._type_info is not None
+                        kwargs[kwarg_node.arg] = kwarg_node.value._type_info.proxy()
                     fn_node = iter_node.func
                     assert isinstance(fn_node, ExtendedAST)
-                    fn = fn_node._type_info.proxy()  # pyright: ignore[reportOptionalMemberAccess]
+                    assert fn_node._type_info is not None
+                    fn = fn_node._type_info.proxy()
                     assert is_api_func(fn)
                     assert fn._codegen is not None
                     bound = fn._signature.bind(*args, **kwargs)