address review comments

Author: isVoid

ast_canopy/ast_canopy/decl.py

@@ -300,7 +300,7 @@ def from_c_obj(cls, c_obj: bindings.Method, parse_entry_point: str):
 
 
 class TemplatedStructMethod(StructMethod):
-    """Struct/class method who's name may include template parameters.
+    """Struct/class method whose name may include template parameters.
 
     Provides utilities for working with the declaration name without
     template arguments.

numbast/src/numbast/function_template.py

@@ -12,7 +12,7 @@
 from ast_canopy.pylibastcanopy import execution_space
 from ast_canopy.decl import FunctionTemplate
 
-from numba import types as nbtypes
+from numba.cuda import types as nbtypes
 from numba.cuda.typing import signature as nb_signature
 from numba.cuda.typing.templates import AbstractTemplate
 from numba.cuda.cudadecl import register, register_global
@@ -38,6 +38,10 @@ def func():
     return func
 
 
+# Registry key: (name, intent_key, shim_writer). name is str; intent_key is tuple
+# or None for intent/overload dispatch; shim_writer is compared by identity so
+# different writer instances get separate entries; make_new_func_obj is the
+# default factory.
 func_obj_registry: dict[tuple[str, tuple | None, object], object] = defaultdict(
     make_new_func_obj
 )

numbast/src/numbast/overload_selection.py

@@ -21,6 +21,9 @@ def _select_templated_overload(
     """
     Select a FunctionTemplate overload for a templated function/method.
 
+    kwds is reserved for future C++-style overload resolution (keyword argument
+    dispatch) and is intentionally unused today.
+
     Today we only select by explicit argument count (visible arity). Keep this
     logic centralized so we can expand it with C++-style overload resolution:
     - filter viable candidates (arity/defaults/variadics, arg_intent visibility),

numbast/tests/test_function_template.py

@@ -14,7 +14,7 @@
 
 
 @pytest.fixture
-def _sample_function_templates():
+def sample_function_templates():
     DATA_FOLDER = os.path.join(os.path.dirname(__file__), "data")
     p = os.path.join(DATA_FOLDER, "sample_function_template.cuh")
     decls = parse_declarations_from_source(p, [p], "sm_80", verbose=True)
@@ -37,8 +37,8 @@ def find_binding(bindings, name):
     )
 
 
-def test_templated_function_overload_selection(_sample_function_templates):
-    func_bindings, shim_writer = _sample_function_templates
+def test_templated_function_overload_selection(sample_function_templates):
+    func_bindings, shim_writer = sample_function_templates
     add = find_binding(func_bindings, "add")
 
     @cuda.jit(link=shim_writer.links())
@@ -55,8 +55,8 @@ def kernel(x, y, z, out):
     np.testing.assert_allclose(out, np.array([3.0, 6.0], dtype=np.float32))
 
 
-def test_templated_function_explicit_specialization(_sample_function_templates):
-    func_bindings, shim_writer = _sample_function_templates
+def test_templated_function_explicit_specialization(sample_function_templates):
+    func_bindings, shim_writer = sample_function_templates
     add = find_binding(func_bindings, "add")
 
     @cuda.jit(link=shim_writer.links())
@@ -76,8 +76,8 @@ def kernel(int_a, int_b, float_a, float_b, out_int, out_float):
     np.testing.assert_allclose(out_float, np.array([4.0], dtype=np.float32))
 
 
-def test_templated_function_default_non_type(_sample_function_templates):
-    func_bindings, shim_writer = _sample_function_templates
+def test_templated_function_default_non_type(sample_function_templates):
+    func_bindings, shim_writer = sample_function_templates
     add_default = find_binding(func_bindings, "add_default")
 
     @cuda.jit(link=shim_writer.links())
@@ -91,8 +91,8 @@ def kernel(inp, out):
     assert out[0] == 17
 
 
-def test_templated_function_default_type(_sample_function_templates):
-    func_bindings, shim_writer = _sample_function_templates
+def test_templated_function_default_type(sample_function_templates):
+    func_bindings, shim_writer = sample_function_templates
     add_default_type = find_binding(func_bindings, "add_default_type")
 
     @cuda.jit(link=shim_writer.links())
@@ -107,9 +107,9 @@ def kernel(inp, out):
 
 
 def test_templated_function_multiple_template_args(
-    _sample_function_templates,
+    sample_function_templates,
 ):
-    func_bindings, shim_writer = _sample_function_templates
+    func_bindings, shim_writer = sample_function_templates
     add_cast = find_binding(func_bindings, "add_cast")
 
     @cuda.jit(link=shim_writer.links())
@@ -125,9 +125,9 @@ def kernel(int_a, float_b, out):
 
 
 def test_templated_function_type_and_non_type(
-    _sample_function_templates,
+    sample_function_templates,
 ):
-    func_bindings, shim_writer = _sample_function_templates
+    func_bindings, shim_writer = sample_function_templates
     add_with_non_type = find_binding(func_bindings, "add_with_non_type")
 
     @cuda.jit(link=shim_writer.links())