mypy fixes

Author: sklam

nbcc/cutile_backend/backend.py

@@ -32,7 +32,7 @@ def entry(
     occupancy=None,
     loc=None,
     ip=None,
-) -> Tile:
+) -> Any:
     """
     from https://github.com/NVIDIA/cuda-tile/blob/8a775693b18303d6c696be6ffd06dadad1b32a8e/python/cuda_tile/dialects/cuda_tile_ops.py#L2C37-L2C44
     """

nbcc/frontend/extra_spy_builtins.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Annotated, cast
+from typing import TYPE_CHECKING, Annotated, cast, Any
 
 from spy.fqn import FQN
 from ..mlir_utils import (
@@ -38,7 +38,7 @@ class W_MLIR_Value(W_Object):
     __spy_storage_category__ = "reference"
 
 
-_type_caches = {}
+_type_caches: dict[str, "W_MLIR_Type"] = {}
 
 
 @MLIR.builtin_type("MLIR_Type")
@@ -102,10 +102,10 @@ def w_opimpl(vm: "SPyVM", *args_w: W_Object) -> W_Object:
 
 @MLIR.builtin_func("MLIR_unpack")
 def w_MLIR_unpack(
-    vm: "SPyVM", w_fn: W_Object, w_idx: W_Object
+    vm: "SPyVM", w_fn: W_Func, w_idx: W_Object
 ) -> W_BuiltinFunc:
 
-    restype = w_fn.w_functype.w_restype
+    restype = cast(W_MLIR_Type, w_fn.w_functype.w_restype)
 
     assert restype.original_name.startswith("multivalues$")
     members = parse_composite_type(restype.original_name)
@@ -115,6 +115,8 @@ def decode_type(str_fqn: str) -> str:
         [enc] = fqn.parts[-1].qualifiers
         return decode_type_name(str(enc))
 
+    assert members is not None
+
     types = list(
         map(lambda x: W_MLIR_Type.w_new(vm, vm.wrap(decode_type(x))), members)
     )
@@ -145,6 +147,7 @@ def w_MLIR_asm(
     vm: "SPyVM", w_asm: W_Str, w_restype: W_Object, w_argtypes: W_Tuple
 ) -> W_BuiltinFunc:
 
+    RESTYPE: Any
     if isinstance(w_restype, W_Tuple):
         innernames = []
         for it_type in w_restype.items_w:

nbcc/mlir_lowering.py

@@ -4,7 +4,13 @@
 from collections import defaultdict
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, Callable, Sequence, cast
+from typing import Any, Callable, Sequence, cast, Coroutine, TYPE_CHECKING
+
+# Type aliases for MLIR types - backends will use their specific ir modules
+IRContext = Any
+IRType = Any
+IRLocation = Any
+IRInsertionPoint = Any
 
 
 from sealir import ase
@@ -70,8 +76,8 @@ class BackendInterface(ABC):
     """
 
     # Class attributes for MLIR context management
-    Location: type  # Set to ir.Location by implementations
-    InsertionPoint: type  # Set to ir.InsertionPoint by implementations
+    Location: Any  # Set to ir.Location by implementations
+    InsertionPoint: Any  # Set to ir.InsertionPoint by implementations
 
     @classmethod
     @abstractmethod
@@ -87,47 +93,47 @@ def run_passes(self, module: Any, transforms: Any) -> Any: ...
     # Type Constants - Properties for clean access pattern
     @property
     @abstractmethod
-    def context(self) -> ir.Context:
+    def context(self) -> "IRContext":
         """MLIR context for creating types and operations."""
 
     @property
     @abstractmethod
-    def i32(self) -> ir.Type:
+    def i32(self) -> "IRType":
         """32-bit integer type."""
 
     @property
     @abstractmethod
-    def i64(self) -> ir.Type:
+    def i64(self) -> "IRType":
         """64-bit integer type."""
 
     @property
     @abstractmethod
-    def f64(self) -> ir.Type:
+    def f64(self) -> "IRType":
         """64-bit float type."""
 
     @property
     @abstractmethod
-    def boolean(self) -> ir.Type:
+    def boolean(self) -> "IRType":
         """Boolean (1-bit integer) type."""
 
     @property
     @abstractmethod
-    def none_type(self) -> ir.Type:
+    def none_type(self) -> "IRType":
         """None/void type representation."""
 
     @property
     @abstractmethod
-    def io_type(self) -> ir.Type:
+    def io_type(self) -> "IRType":
         """IO token type for sequencing."""
 
     @property
     @abstractmethod
-    def llvm_ptr(self) -> ir.Type:
+    def llvm_ptr(self) -> "IRType":
         """LLVM pointer type for memory operations."""
 
     # Core Methods
     @abstractmethod
-    def lower_type(self, ty) -> tuple[ir.Type, ...]:
+    def lower_type(self, ty) -> tuple["IRType", ...]:
         """Convert SealIR types to backend IR types.
 
         Returns a tuple of MLIR types. For single types, returns (type,).
@@ -136,7 +142,7 @@ def lower_type(self, ty) -> tuple[ir.Type, ...]:
         """
 
     @abstractmethod
-    def get_ll_type(self, expr, mdmap) -> ir.Type | None:
+    def get_ll_type(self, expr, mdmap) -> "IRType | None":
         """Get backend type for expression with metadata."""
 
     @abstractmethod
@@ -280,7 +286,7 @@ def lower(self, root: rg.Func) -> Any:
         """
         context = self.be.context
         self.loc = loc = self.be.Location.name(
-            f"{self}.lower()", context=context
+            f"{self.__class__.__name__}.lower()", context=context
         )
         module = self.module
 

pyproject.toml

@@ -32,7 +32,7 @@ mypy_path = ["deps/sealir", "deps/spy"]
 python_version = "3.12"
 
 [[tool.mypy.overrides]]
-module = ["mlir.*", "sealir.*", "spy.*"]
+module = ["mlir.*", "sealir.*", "spy.*", "cuda_tile.*"]
 ignore_errors = true
 ignore_missing_imports = true