Call generic_visit in all paths and bail out early in all visit methods

src/pytest_run_parallel/thread_unsafe_detection.py

@@ -138,18 +138,13 @@ def _visit_name_call(self, node):
             self._recursive_analyze_name(node)
 
     def visit_Call(self, node):
-        if self.thread_unsafe:
-            return
-
         if isinstance(node.func, ast.Attribute):
             self._visit_attribute_call(node.func)
         elif isinstance(node.func, ast.Name):
             self._visit_name_call(node.func)
+        self.generic_visit(node)
 
     def visit_Assign(self, node):
-        if self.thread_unsafe:
-            return
-
         if len(node.targets) == 1:
             name_node = node.targets[0]
             value_node = node.value
@@ -159,8 +154,13 @@ def visit_Assign(self, node):
                     f"calls thread-unsafe function: f{name_node} "
                     "(inferred via func.__thread_safe__ == False)"
                 )
-            else:
-                self.generic_visit(node)
+
+        self.generic_visit(node)
+
+    def visit(self, node):
+        if self.thread_unsafe:
+            return
+        return super().visit(node)
 
 
 def _identify_thread_unsafe_nodes(fn, skip_set, level=0):

tests/test_thread_unsafe_detection.py

@@ -367,3 +367,40 @@ def test_chained_attribute_thread_safe_assignment(num_parallel_threads):
             "*::test_chained_attribute_thread_safe_assignment PASSED*",
         ]
     )
+
+
+def test_wrapped_function_call(pytester):
+    pytester.makepyfile("""
+    import pytest
+
+    def wrapper(x):
+        return x
+
+    def test_wrapped_function_call(num_parallel_threads):
+        wrapper(pytest.warns())
+        assert num_parallel_threads == 1
+    """)
+
+    result = pytester.runpytest("--parallel-threads=10", "-v")
+    result.stdout.fnmatch_lines(
+        [
+            "*::test_wrapped_function_call PASSED*",
+        ]
+    )
+
+
+def test_thread_unsafe_function_call_in_assignment(pytester):
+    pytester.makepyfile("""
+    import pytest
+
+    def test_thread_unsafe_function_call_in_assignment(num_parallel_threads):
+        x = y = pytest.warns()
+        assert num_parallel_threads == 1
+    """)
+
+    result = pytester.runpytest("--parallel-threads=10", "-v")
+    result.stdout.fnmatch_lines(
+        [
+            "*::test_thread_unsafe_function_call_in_assignment PASSED*",
+        ]
+    )