remove the add_code_object public entry and update tests

Author: ksimpson-work

cuda_core/cuda/core/experimental/_linker.py

@@ -92,14 +92,12 @@ def __init__(self, *object_codes : ObjectCode, options: LinkerOptions = None):
         self._handle = nvjitlink.create(len(options.formatted_options), options.formatted_options)
 
         if object_codes is not None:
-            if object_codes.__iter__:
-                for code in object_codes:
-                    self.add_code_object(code)
-            else:
-                self.add_code_object(object_codes)
+            for code in object_codes:
+                assert isinstance(code, ObjectCode)
+                self._add_code_object(code)
 
 
-    def add_code_object(self, object_code: ObjectCode):
+    def _add_code_object(self, object_code: ObjectCode):
         data = object_code._module
         assert isinstance(data, bytes)
         nvjitlink.add_data(

cuda_core/tests/test_linker.py

@@ -50,66 +50,49 @@ def compile_ltoir_functions(init_cuda):
     LinkerOptions(arch=ARCH, no_cache=True)
 ])
 def test_linker_init(compile_ptx_functions, options):
-    linker = Linker(options, object_codes=compile_ptx_functions)
+    linker = Linker(*compile_ptx_functions, options=options)
     object_code = linker.link("cubin")
     assert isinstance(object_code, ObjectCode)
 
 
 def test_linker_init_invalid_arch():
     options = LinkerOptions(arch=None)
-    with pytest.raises(ValueError):
+    with pytest.raises(TypeError):
         Linker(options)
 
 
-def test_linker_add_code_object(compile_ptx_functions):
-    options = LinkerOptions(arch=ARCH)
-    linker = Linker(options)
-    for functions in compile_ptx_functions:
-        linker.add_code_object(functions)
-
-
 def test_linker_link_ptx(compile_ltoir_functions):
     options = LinkerOptions(arch=ARCH, link_time_optimization=True, ptx=True)
-    linker = Linker(options)
-    for functions in compile_ltoir_functions:
-        linker.add_code_object(functions)
+    linker = Linker(*compile_ltoir_functions, options = options)
     linked_code = linker.link("ptx")
     assert isinstance(linked_code, ObjectCode)
 
 
 def test_linker_link_cubin(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
-    linker = Linker(options)
-    for functions in compile_ptx_functions:
-        linker.add_code_object(functions)
+    linker = Linker(*compile_ptx_functions, options=options)
     linked_code = linker.link("cubin")
     assert isinstance(linked_code, ObjectCode)
 
 
 def test_linker_link_invalid_target_type(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
-    linker = Linker(options)
-    for functions in compile_ptx_functions:
-        linker.add_code_object(functions)
+    linker = Linker(*compile_ptx_functions, options=options)
     with pytest.raises(ValueError):
         linker.link("invalid_target")
 
 
 def test_linker_get_error_log(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
-    linker = Linker(options)
-    for functions in compile_ptx_functions:
-        linker.add_code_object(functions)
+    linker = Linker(*compile_ptx_functions, options=options)
     linker.link("cubin")
     log = linker.get_error_log()
     assert isinstance(log, str)
 
 
 def test_linker_get_info_log(compile_ptx_functions):
     options = LinkerOptions(arch=ARCH)
-    linker = Linker(options)
-    for functions in compile_ptx_functions:
-        linker.add_code_object(functions)
+    linker = Linker(*compile_ptx_functions, options=options)
     linker.link("cubin")
     log = linker.get_info_log()
     assert isinstance(log, str)