Refactored solution by review requests

testplan/testing/base.py

@@ -693,7 +693,9 @@ def _run_resource_hook(
         )
 
         case_result = self.cfg.result(
-            stdout_style=self.stdout_style, _scratch=self.scratch
+            stdout_style=self.stdout_style,
+            _scratch=self.scratch,
+            _collect_code_context=self.collect_code_context,
         )
         runtime_env = self._get_runtime_environment(
             testcase_name=hook_name,
@@ -907,7 +909,7 @@ def collect_code_context(self) -> bool:
         Collecting the file path, line number and code context of the assertions
         if enabled.
         """
-        return self.cfg.collect_code_context
+        return getattr(self.cfg, "collect_code_context", False)
 
 
 class ProcessRunnerTestConfig(TestConfig):

testplan/testing/multitest/base.py

@@ -30,7 +30,7 @@
 )
 from testplan.testing.multitest import suite as mtest_suite
 from testplan.testing.multitest.entries import base as entries_base
-from testplan.testing.result import report_target, collect_code_context
+from testplan.testing.result import report_target
 from testplan.testing.multitest.suite import (
     get_suite_metadata,
     get_testcase_metadata,
@@ -1045,7 +1045,9 @@ def _run_suite_related(self, testsuite, method_name):
 
         method_report = self._suite_related_report(method_name)
         case_result = self.cfg.result(
-            stdout_style=self.stdout_style, _scratch=self._scratch
+            stdout_style=self.stdout_style,
+            _scratch=self._scratch,
+            _collect_code_context=self.collect_code_context,
         )
 
         resources = self._get_runtime_environment(
@@ -1137,7 +1139,9 @@ def _run_testcase(
             testcase
         )
         case_result: result.Result = self.cfg.result(
-            stdout_style=self.stdout_style, _scratch=self.scratch
+            stdout_style=self.stdout_style,
+            _scratch=self.scratch,
+            _collect_code_context=self.collect_code_context,
         )
 
         # as the runtime info currently has only testcase name we create it here
@@ -1148,18 +1152,15 @@ def _run_testcase(
             testcase_name=testcase.name, testcase_report=testcase_report
         )
 
-        if getattr(self.cfg, "collect_code_context", False):
-            if self.cfg.testcase_report_target:
-                testcase = report_target(
-                    func=testcase,
-                    ref_func=getattr(
-                        testsuite,
-                        getattr(testcase, "_parametrization_template", ""),
-                        None,
-                    ),
-                )
-
-            testcase = collect_code_context(func=testcase)
+        if self.cfg.testcase_report_target and self.collect_code_context:
+            testcase = report_target(
+                func=testcase,
+                ref_func=getattr(
+                    testsuite,
+                    getattr(testcase, "_parametrization_template", ""),
+                    None,
+                ),
+            )
 
         # specially handle skipped testcases
         if hasattr(testcase, "__should_skip__"):

testplan/testing/result.py

@@ -93,7 +93,7 @@ def __exit__(self, exc_type, exc_value, tb):
             description=self.description,
         )
 
-        if getattr(assertion_state, "collect_code_context", False):
+        if self.result._collect_code_context:
             with MOD_LOCK:
                 # TODO: see https://github.com/python/cpython/commit/85cf1d514b84dc9a4bcb40e20a12e1d82ff19f20
                 caller_frame = inspect.stack()[1]
@@ -115,19 +115,6 @@ def __exit__(self, exc_type, exc_value, tb):
 assertion_state = threading.local()
 
 
-def collect_code_context(func: Callable) -> Callable:
-    """
-    Sets the decorated function to collect code context
-    """
-
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        assertion_state.collect_code_context = True
-        func(*args, **kwargs)
-
-    return wrapper
-
-
 def report_target(func: Callable, ref_func: Callable = None) -> Callable:
     """
     Sets the decorated function's filepath and line-range in assertion state.
@@ -183,7 +170,21 @@ def wrapper(result, *args, **kwargs):
             if not top_assertion:
                 return entry
 
-            if getattr(assertion_state, "collect_code_context", False):
+            if custom_style is not None:
+                if not isinstance(custom_style, dict):
+                    raise TypeError(
+                        "Use `dict[str, str]` to specify custom CSS style"
+                    )
+                entry.custom_style = custom_style
+
+            assert isinstance(result, AssertionNamespace) or isinstance(
+                result, Result
+            ), "Incorrect usage of assertion decorator"
+
+            if isinstance(result, AssertionNamespace):
+                result = result.result
+
+            if result._collect_code_context:
                 with MOD_LOCK:
                     call_stack = inspect.stack()
                     try:
@@ -207,20 +208,6 @@ def wrapper(result, *args, **kwargs):
                         del frame
                         del call_stack
 
-            if custom_style is not None:
-                if not isinstance(custom_style, dict):
-                    raise TypeError(
-                        "Use `dict[str, str]` to specify custom CSS style"
-                    )
-                entry.custom_style = custom_style
-
-            assert isinstance(result, AssertionNamespace) or isinstance(
-                result, Result
-            ), "Incorrect usage of assertion decorator"
-
-            if isinstance(result, AssertionNamespace):
-                result = result.result
-
             if not dryrun:
                 result.entries.append(entry)
 
@@ -1390,7 +1377,7 @@ def __exit__(self, exc_type, exc_value, traceback):
             return False
         super().__exit__(exc_type, exc_value, traceback)
 
-        if getattr(assertion_state, "collect_code_context", False):
+        if self.result._collect_code_context:
             with MOD_LOCK:
                 # TODO: see https://github.com/python/cpython/commit/85cf1d514b84dc9a4bcb40e20a12e1d82ff19f20
                 # XXX: do we have concrete ideas about thread-safety here?
@@ -1559,6 +1546,7 @@ def __init__(
         _num_passing=defaults.SUMMARY_NUM_PASSING,
         _num_failing=defaults.SUMMARY_NUM_FAILING,
         _scratch=None,
+        _collect_code_context=False,
     ):
 
         self.entries = []
@@ -1580,6 +1568,7 @@ def __init__(
         self._num_passing = _num_passing
         self._num_failing = _num_failing
         self._scratch = _scratch
+        self._collect_code_context = _collect_code_context
 
     def subresult(self):
         """Subresult object to append/prepend assertions on another."""

tests/unit/testplan/testing/multitest/test_result.py

@@ -54,15 +54,14 @@ def test_group_no_marking():
     assert result_entry.code_context == None
 
 
-@result_mod.collect_code_context
 def test_group_marking():
     """
     Tests, at result object level, if marking works as expected.
     """
-    result = result_mod.Result()
+    result = result_mod.Result(_collect_code_context=True)
     result.equal(1, 1)
     result_entry = result.entries.pop()
-    code_context = get_code_context(test_group_marking, 6)
+    code_context = get_code_context(test_group_marking, 5)
     assert result_entry.line_no == code_context[0]
     assert result_entry.code_context == code_context[1]