standalone-tsm-context - get_context_with_sites

src/nitsm/codemoduleapi.py

@@ -42,13 +42,20 @@ def __call__(self, *args, **kwargs):
             if inspect.isclass(argument[1]):  # class method check
                 argument = next(arguments_iter)  # move to second argument
         except StopIteration:
-            raise TypeError(
-                (
-                    "The number of arguments to the code module is less than expected. It must "
-                    "accept as it's first argument the Semiconductor Module context passed from "
-                    "TestStand or another code module.",
+            # Fallback: look for tsm context in kwargs by name
+            for key in ("context", "tsm_context", "SemiconductorModuleContext"):
+                if key in bound_arguments.arguments:
+                    argument = (key, bound_arguments.arguments[key])
+                    break
+            else:
+                raise TypeError(
+                    (
+                        "The number of arguments to the code module is less than expected. It must "
+                        "accept as it's first argument the Semiconductor Module context passed from "
+                        "TestStand or another code module or it should be passed as a keyword argument,"
+                        " with keywords - 'context', 'tsm_context', or 'SemiconductorModuleContext'.",
+                    )
                 )
-            )
 
         # attempt to wrap argument in a SemiconductorModuleContext object
         argument_name, argument_value = argument

src/nitsm/tsmcontext.py

@@ -286,6 +286,40 @@ def global_data_exists(self, data_id: str) -> bool:
         """
         return self._context.GlobalDataExists(data_id)
 
+    def get_semiconductor_module_context_with_sites(self,site_numbers: _Sequence[int]) -> "SemiconductorModuleContext":
+        """Returns a Semiconductor Module context object which holds information and resources
+        specific to the site_numbers mentioned.
+
+        Args:
+            site_numbers: A sequence of site numbers for which the resources should be used.
+
+        Returns:
+            SemiconductorModuleContext object with resources specific to the site_numbers.
+        """
+        # Normalize and validate the site numbers before passing them to the underlying
+        # Semiconductor Module context. This ensures predictable behavior for invalid or
+        # empty site lists.
+        if not site_numbers:
+            raise ValueError("site_numbers must contain at least one site number.")
+        normalized_site_numbers = []
+
+        for site in site_numbers:
+            if not isinstance(site, int) or site < 0:
+                raise ValueError("Each site number must be a non-negative integer.")
+            normalized_site_numbers.append(site)
+
+        # We have to use DumbDispatch here because pywin32 fails to recognize
+        # ISemiconductorModuleContext as deriving from IDispatch; most likely because it isn't
+        # natively supported. So, we fetch it as IUnknown instead.
+        dumb_context = win32com.client.dynamic.DumbDispatch(self._context)
+
+        # Retrieve the SemiconductorModuleContext COM pointer for the selected sites
+        context_with_sites = dumb_context.GetSemiconductorModuleContextWithSites(normalized_site_numbers)
+
+        # Wrap the cleaned IUnknown pointer in the Python SemiconductorModuleContext wrapper
+        return SemiconductorModuleContext(context_with_sites)
+
+
     # NI-Digital
 
     def get_all_nidigital_instrument_names(self) -> "_StringTuple":

tests/test_codemoduleapi.py

@@ -1,5 +1,6 @@
 import re
 import pytest
+from unittest.mock import MagicMock
 from nitsm.codemoduleapi import code_module, SemiconductorModuleContext
 
 
@@ -56,7 +57,8 @@ def test_invalid_number_of_positional_arguments(self):
         assert e.value.args[0] == (
             "The number of arguments to the code module is less than expected. It must "
             "accept as it's first argument the Semiconductor Module context passed from "
-            "TestStand or another code module.",
+            "TestStand or another code module or it should be passed as a keyword argument,"
+            " with keywords - 'context', 'tsm_context', or 'SemiconductorModuleContext'.",
         )
 
     # noinspection PyUnusedLocal
@@ -84,6 +86,25 @@ def test_positional_arguments_after_tsm_context(
     ):
         assert isinstance(standalone_tsm_context_com_object, SemiconductorModuleContext)
 
+
+    def test_tsm_context_via_kwargs(self):
+        raw_com_context = MagicMock(name="RawTSMContext")
+
+        received_context = None
+
+        @code_module
+        def my_module(tsm_context=None):
+            nonlocal received_context
+            received_context = tsm_context
+            return "converted"
+
+        result = my_module(tsm_context=raw_com_context)
+
+        # The function returned normally
+        assert result == "converted"
+
+        # Check that the wrapper converted raw COM object → SemiconductorModuleContext
+        assert isinstance(received_context, SemiconductorModuleContext)
 
 @pytest.mark.pin_map("codemoduleapi.pinmap")
 @code_module