Factor out common logic between num_arguments and arguments_info properties

Also parametrize test to check arguments_info to check with all int,
and all short arguments.

Author: oleksandr-pavlyk

cuda_core/cuda/core/experimental/_module.py

@@ -216,34 +216,31 @@ def attributes(self) -> KernelAttributes:
             self._attributes = KernelAttributes._init(self._handle)
         return self._attributes
 
-    @property
-    def num_arguments(self) -> int:
-        """int : The number of arguments of this function"""
+    def _get_arguments_info(self, param_info=False) -> tuple[int, list[tuple[int, int]]]:
         attr_impl = self.attributes
         if attr_impl._backend_version != "new":
             raise NotImplementedError("New backend is required")
         arg_pos = 0
+        param_info_data = []
         while True:
             result = attr_impl._loader["paraminfo"](self._handle, arg_pos)
             if result[0] != driver.CUresult.CUDA_SUCCESS:
                 break
+            if param_info:
+                param_info_data.append(result[1:3])
             arg_pos = arg_pos + 1
-        return arg_pos
+        return arg_pos, param_info_data
+
+    @property
+    def num_arguments(self) -> int:
+        """int : The number of arguments of this function"""
+        num_args, _ = self._get_arguments_info()
+        return num_args
 
     @property
     def arguments_info(self) -> list[tuple[int, int]]:
-        """tuple[tuple[int, int]]: (offset, size) for each argument of this function"""
-        attr_impl = self.attributes
-        if attr_impl._backend_version != "new":
-            raise NotImplementedError("New backend is required")
-        arg_pos = 0
-        param_info = []
-        while True:
-            result = attr_impl._loader["paraminfo"](self._handle, arg_pos)
-            if result[0] != driver.CUresult.CUDA_SUCCESS:
-                break
-            param_info.append((result[1], result[2]))
-            arg_pos = arg_pos + 1
+        """list[tuple[int, int]]: (offset, size) for each argument of this function"""
+        _, param_info = self._get_arguments_info(param_info=True)
         return param_info
 
     # TODO: implement from_handle()

cuda_core/tests/test_module.py

@@ -8,6 +8,7 @@
 
 
 import warnings
+import ctypes
 
 import pytest
 
@@ -165,7 +166,6 @@ def test_object_code_handle(get_saxpy_object_code):
 
 
 def test_saxpy_arguments(get_saxpy_kernel):
-    import ctypes
     krn, _ = get_saxpy_kernel
 
     assert krn.num_arguments == 5
@@ -190,8 +190,9 @@ class ExpectedStruct(ctypes.Structure):
 
 
 @pytest.mark.parametrize("nargs", [0, 1, 2, 3, 16])
-def test_num_arguments(init_cuda, nargs):
-    args_str = ", ".join([f"int p_{i}" for i in range(nargs)])
+@pytest.mark.parametrize("c_type_name,c_type", [("int", ctypes.c_int), ("short", ctypes.c_short)], ids=["int", "short"])
+def test_num_arguments(init_cuda, nargs, c_type_name, c_type):
+    args_str = ", ".join([f"{c_type_name} p_{i}" for i in range(nargs)])
     src = f"__global__ void foo{nargs}({args_str}) {{ }}"
     prog = Program(src, code_type="c++")
     mod = prog.compile(
@@ -201,10 +202,9 @@ def test_num_arguments(init_cuda, nargs):
     krn = mod.get_kernel(f"foo{nargs}")
     assert krn.num_arguments == nargs
 
-    import ctypes
     class ExpectedStruct(ctypes.Structure):
         _fields_ = [
-            (f'arg_{i}', ctypes.c_int) for i in range(nargs)
+            (f'arg_{i}', c_type) for i in range(nargs)
         ]
     members = tuple(getattr(ExpectedStruct, f"arg_{i}") for i in range(nargs))
     expected = [(m.offset, m.size) for m in members]