inducer · inducer · Jan 10, 2025 · Dec 27, 2024 · Dec 27, 2024 · Jan 9, 2025
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -57,7 +57,7 @@ jobs:
                 . ./ci-support-v0
 
                 build_py_project_in_venv
-                python -m mypy codepy
+                python -m mypy codepy test examples
 
     docs:
         name: Documentation

diff --git a/codepy/bpl.py b/codepy/bpl.py
@@ -1,38 +1,47 @@
 """Convenience interface for using CodePy with Boost.Python."""
 
+from collections.abc import Callable, Iterable
 from dataclasses import replace
+from types import ModuleType
+from typing import 
8000
Any
+
+from cgen import FunctionBody, Generable, Module, Struct
+
+from codepy.toolchain import Toolchain
 
 
 class BoostPythonModule:
-    def __init__(self, name="module", max_arity=None,
-            use_private_namespace=True):
+    def __init__(self,
+                 name: str = "module",
+                 max_arity: int | str | None = None,
+                 use_private_namespace: bool = True) -> None:
         self.name = name
-        self.preamble = []
-        self.mod_body = []
-        self.init_body = []
+        self.max_arity = max_arity
+        self.use_private_namespace = use_private_namespace
 
+        self.preamble: list[Generable] = []
+        self.mod_body: list[Generable] = []
+        self.init_body: list[Generable] = []
         self.has_codepy_include = False
         self.has_raw_function_include = False
-        self.max_arity = max_arity
-        self.use_private_namespace = use_private_namespace
 
-    def add_to_init(self, body):
+    def add_to_init(self, body: Iterable[Generable]) -> None:
         """Add the blocks or statements contained in the iterable *body* to the
         module initialization function.
         """
         self.init_body.extend(body)
 
-    def add_to_preamble(self, pa):
+    def add_to_preamble(self, pa: Iterable[Generable]) -> None:
         self.preamble.extend(pa)
 
-    def add_to_module(self, body):
+    def add_to_module(self, body: Iterable[Generable]) -> None:
         """Add the :class:`cgen.Generable` instances in the iterable
         *body* to the body of the module *self*.
         """
 
         self.mod_body.extend(body)
 
-    def add_codepy_include(self):
+    def add_codepy_include(self) -> None:
         if self.has_codepy_include:
             return
 
@@ -43,7 +52,7 @@ def add_codepy_include(self):
             ])
         self.has_codepy_include = True
 
-    def add_raw_function_include(self):
+    def add_raw_function_include(self) -> None:
         if self.has_raw_function_include:
             return
 
@@ -54,7 +63,7 @@ def add_raw_function_include(self):
             ])
         self.has_raw_function_include = True
 
-    def expose_vector_type(self, name, py_name=None):
+    def expose_vector_type(self, name: str, py_name: str | None = None) -> None:
         self.add_codepy_include()
 
         if py_name is None:
@@ -71,34 +80,39 @@ def expose_vector_type(self, name, py_name=None):
                     ".def(codepy::no_compare_indexing_suite<cl>())"),
                 ]))
 
-    def add_function(self, func):
+    def add_function(self, func: FunctionBody) -> None:
         """Add a function to be exposed. *func* is expected to be a
         :class:`cgen.FunctionBody`.
         """
+        from cgen import Statement
 
         self.mod_body.append(func)
-        from cgen import Statement
         self.init_body.append(
                 Statement(
                     'boost::python::def("{}", &{})'.format(
                         func.fdecl.name, func.fdecl.name)))
 
-    def add_raw_function(self, func):
+    def add_raw_function(self, func: FunctionBody) -> None:
         """Add a function to be exposed using boost::python::raw_function.
         *func* is expected to be a :class:`cgen.FunctionBody`.
         """
-        self.mod_body.append(func)
         from cgen import Statement
+
+        self.mod_body.append(func)
         self.add_raw_function_include()
+
         raw_function = f"boost::python::raw_function(&{func.fdecl.name})"
         self.init_body.append(
             Statement(
                 'boost::python::def("{}", {})'.format(
                     func.fdecl.name, raw_function)))
 
-    def add_struct(self,
-            struct, py_name=None, py_member_name_transform=lambda x: x,
-            by_value_members=None):
+    def add_struct(
+            self,
+            struct: Struct,
+            py_name: str | None = None,
+            py_member_name_transform: Callable[[str], str] = lambda x: x,
+            by_value_members: set[str] | None = None) -> None:
         if by_value_members is None:
             by_value_members = set()
 
@@ -111,6 +125,11 @@ def add_struct(self,
 
         member_defs = []
         for f in struct.fields:
+            if not hasattr(f, "name"):
+                raise TypeError(
+                    f"Invalid type {type(f)} of struct field. Only named fields "
+                    "are supported for code generation")
+
             py_f_name = py_member_name_transform(f.name)
             tp_lines, _ = f.get_decl_pair()
             if f.name in by_value_members or tp_lines[0].startswith("numpy_"):
@@ -131,20 +150,20 @@ def add_struct(self,
                         py_name, "".join(member_defs))),
                 ]))
 
-    def generate(self):
+    def generate(self) -> Module:
         """Generate (i.e. yield) the source code of the
         module line-by-line.
         """
 
-        from cgen import Block, Define, Include, Line, Module, PrivateNamespace
+        from cgen import Block, Define, Include, Line, PrivateNamespace
 
-        body = []
+        body: list[Generable] = []
 
         if self.max_arity is not None:
-            body.append(Define("BOOST_PYTHON_MAX_ARITY", self.max_arity))
+            body.append(Define("BOOST_PYTHON_MAX_ARITY", str(self.max_arity)))
 
         if self.use_private_namespace:
-            mod_body = [PrivateNamespace(self.mod_body)]
+            mod_body: list[Generable] = [PrivateNamespace(self.mod_body)]
         else:
             mod_body = self.mod_body
 
@@ -160,7 +179,7 @@ def generate(self):
 
         return Module(body)
 
-    def compile(self, toolchain, **kwargs):
+    def compile(self, toolchain: Toolchain, **kwargs: Any) -> ModuleType:
         """Return the extension module generated from the code described
         by *self*. If necessary, build the code using *toolchain* with
         :func:`codepy.jit.extension_from_string`. Any keyword arguments
@@ -173,5 +192,8 @@ def compile(self, toolchain, **kwargs):
         add_boost_python(toolchain)
 
         from codepy.jit import extension_from_string
-        return extension_from_string(toolchain, self.name,
+
+        return extension_from_string(
+                toolchain,
+                self.name,
                 "{}\n".format(self.generate()), **kwargs)
diff --git a/codepy/cuda.py b/codepy/cuda.py
@@ -1,13 +1,20 @@
+"""Convenience interface for using CodePy with CUDA"""
+
+from collections.abc import Iterable
 from dataclasses import replace
+from types import ModuleType
+from typing import Any
 
 import cgen
 
-
-"""Convenience interface for using CodePy with CUDA"""
+from codepy.bpl import BoostPythonModule
+from codepy.toolchain import NVCCToolchain, Toolchain
 
 
 class CudaModule:
-    def __init__(self, boost_module, name="module"):
+    def __init__(self,
+                 boost_module: BoostPythonModule,
+                 name: str = "module") -> None:
         """*boost_module* is a codepy.BoostPythonModule containing host code
         which calls CUDA code.
         Current limitations of nvcc preclude compiling anything which
@@ -19,21 +26,21 @@ def __init__(self, boost_module, name="module"):
         """
         self.name = name
 
-        self.preamble = []
-        self.body = []
+        self.preamble: list[cgen.Generable] = []
+        self.body: list[cgen.Generable] = []
         self.boost_module = boost_module
         self.boost_module.add_to_preamble([cgen.Include("cuda.h")])
 
-    def add_to_preamble(self, pa):
+    def add_to_preamble(self, pa: Iterable[cgen.Generable]) -> None:
         self.preamble.extend(pa)
 
-    def add_to_module(self, body):
+    def add_to_module(self, body: Iterable[cgen.Generable]) -> None:
         """Add the :class:`cgen.Generable` instances in the iterable
         *body* to the body of the module *self*.
         """
         self.body.extend(body)
 
-    def add_function(self, func):
+    def add_function(self, func: cgen.FunctionBody) -> None:
         """Add a function to be exposed to code in the BoostPythonModule.
         *func* is expected to be a :class:`cgen.FunctionBody`.
         Additionally, *func* must be a host callable function,
@@ -42,16 +49,20 @@ def add_function(self, func):
         self.boost_module.add_to_preamble([func.fdecl])
         self.body.append(func)
 
-    def generate(self):
+    def generate(self) -> cgen.Module:
         """Generate (i.e. yield) the source code of the
         module line-by-line.
         """
         body = []
         body += [*self.preamble, cgen.Line(), *self.body]
         return cgen.Module(body)
 
-    def compile(self, host_toolchain, nvcc_toolchain,
-            host_kwargs=None, nvcc_kwargs=None, **kwargs):
+    def compile(self,
+                host_toolchain: Toolchain,
+                nvcc_toolchain: NVCCToolchain,
+                host_kwargs: dict[str, Any] | None = None,
+                nvcc_kwargs: dict[str, Any] | None = None,
+                **kwargs: Any) -> ModuleType:
         """Return the extension module generated from the code described
         by *self*. If necessary, build the code using *toolchain* with
         :func:`codepy.jit.extension_from_string`. Any keyword arguments
@@ -104,11 +115,13 @@ def compile(self, host_toolchain, nvcc_toolchain,
         else:
             import os.path
 
+            ext = getattr(host_toolchain, "so_ext", ".so")
             destination_base, _ = os.path.split(host_object)
-            module_path = os.path.join(destination_base, mod_name
-                                       + host_toolchain.so_ext)
+            module_path = os.path.join(destination_base, f"{mod_name}{ext}")
+
+            from codepy.tools import load_dynamic
+
             try:
-                from codepy.tools import load_dynamic
                 return load_dynamic(mod_name, module_path)
             except Exception:
                 return link_extension(host_toolchain,