Refactor plugin into multiple files

Author: lysnikolaou

src/pytest_run_parallel/cpu_detection.py

@@ -0,0 +1,36 @@
+def get_logical_cpus():
+    try:
+        import psutil
+    except ImportError:
+        pass
+    else:
+        process = psutil.Process()
+        try:
+            cpu_cores = process.cpu_affinity()
+            return len(cpu_cores)
+        except AttributeError:
+            cpu_cores = psutil.cpu_count()
+            if cpu_cores is not None:
+                return cpu_cores
+
+    try:
+        from os import process_cpu_count
+    except ImportError:
+        pass
+    else:
+        cpu_cores = process_cpu_count()
+        if cpu_cores is not None:
+            return cpu_cores
+
+    try:
+        from os import sched_getaffinity
+    except ImportError:
+        pass
+    else:
+        cpu_cores = sched_getaffinity(0)
+        if cpu_cores is not None:
+            return len(cpu_cores)
+
+    from os import cpu_count
+
+    return cpu_count()

src/pytest_run_parallel/fixtures.py

@@ -0,0 +1,24 @@
+import pytest
+
+from pytest_run_parallel.thread_comparator import ThreadComparator
+from pytest_run_parallel.utils import get_num_workers
+
+
+@pytest.fixture
+def num_parallel_threads(request):
+    return get_num_workers(request.config, request.node)
+
+
+@pytest.fixture
+def num_iterations(request):
+    node = request.node
+    n_iterations = request.config.option.iterations
+    m = node.get_closest_marker("iterations")
+    if m is not None:
+        n_iterations = int(m.args[0])
+    return n_iterations
+
+
+@pytest.fixture
+def thread_comp(num_parallel_threads):
+    return ThreadComparator(num_parallel_threads)

src/pytest_run_parallel/plugin.py renamed to src/pytest_run_parallel/hooks.py

@@ -1,18 +1,14 @@
-import functools
 import os
-import sys
-import threading
 import warnings
 
-import _pytest.outcomes
 import pytest
 
-from pytest_run_parallel.utils import (
-    ThreadComparator,
-    get_configured_num_workers,
-    get_num_workers,
+from pytest_run_parallel.thread_unsafe_detection import (
+    THREAD_UNSAFE_FIXTURES,
     identify_thread_unsafe_nodes,
 )
+from pytest_run_parallel.thread_wrapper import wrap_function_parallel
+from pytest_run_parallel.utils import get_configured_num_workers, get_num_workers
 
 
 def pytest_addoption(parser):
@@ -72,81 +68,6 @@ def pytest_configure(config):
     )
 
 
-def wrap_function_parallel(fn, n_workers, n_iterations):
-    @functools.wraps(fn)
-    def inner(*args, **kwargs):
-        errors = []
-        skip = None
-        failed = None
-        barrier = threading.Barrier(n_workers)
-        original_switch = sys.getswitchinterval()
-        new_switch = 1e-6
-        for _ in range(3):
-            try:
-                sys.setswitchinterval(new_switch)
-                break
-            except ValueError:
-                new_switch *= 10
-        else:
-            sys.setswitchinterval(original_switch)
-
-        try:
-
-            def closure(*args, **kwargs):
-                for _ in range(n_iterations):
-                    barrier.wait()
-                    try:
-                        fn(*args, **kwargs)
-                    except Warning:
-                        pass
-                    except Exception as e:
-                        errors.append(e)
-                    except _pytest.outcomes.Skipped as s:
-                        nonlocal skip
-                        skip = s.msg
-                    except _pytest.outcomes.Failed as f:
-                        nonlocal failed
-                        failed = f
-
-            workers = []
-            for _ in range(0, n_workers):
-                worker_kwargs = kwargs
-                workers.append(
-                    threading.Thread(target=closure, args=args, kwargs=worker_kwargs)
-                )
-
-            num_completed = 0
-            try:
-                for worker in workers:
-                    worker.start()
-                    num_completed += 1
-            finally:
-                if num_completed < len(workers):
-                    barrier.abort()
-
-            for worker in workers:
-                worker.join()
-
-        finally:
-            sys.setswitchinterval(original_switch)
-
-        if skip is not None:
-            pytest.skip(skip)
-        elif failed is not None:
-            raise failed
-        elif errors:
-            raise errors[0]
-
-    return inner
-
-
-_thread_unsafe_fixtures = {
-    "capsys",
-    "monkeypatch",
-    "recwarn",
-}
-
-
 @pytest.hookimpl(trylast=True)
 def pytest_itemcollected(item):
     n_workers = get_num_workers(item.config, item)
@@ -207,7 +128,7 @@ def pytest_itemcollected(item):
             else:
                 item.add_marker(pytest.mark.parallel_threads(1))
 
-    unsafe_fixtures = _thread_unsafe_fixtures | set(
+    unsafe_fixtures = THREAD_UNSAFE_FIXTURES | set(
         item.config.getini("thread_unsafe_fixtures")
     )
 
@@ -307,23 +228,3 @@ def pytest_terminal_summary(terminalreporter, exitstatus, config):
             )
     if n_workers > 1 and num_serial == 0:
         terminalreporter.line("All tests were run in parallel! 🎉")
-
-
-@pytest.fixture
-def num_parallel_threads(request):
-    return get_num_workers(request.config, request.node)
-
-
-@pytest.fixture
-def num_iterations(request):
-    node = request.node
-    n_iterations = request.config.option.iterations
-    m = node.get_closest_marker("iterations")
-    if m is not None:
-        n_iterations = int(m.args[0])
-    return n_iterations
-
-
-@pytest.fixture
-def thread_comp(num_parallel_threads):
-    return ThreadComparator(num_parallel_threads)

src/pytest_run_parallel/thread_comparator.py

@@ -0,0 +1,73 @@
+import threading
+import types
+
+try:
+    import numpy as np
+
+    numpy_available = True
+except ImportError:
+    numpy_available = False
+
+
+class ThreadComparator:
+    def __init__(self, n_threads):
+        self._barrier = threading.Barrier(n_threads)
+        self._reset_evt = threading.Event()
+        self._entry_barrier = threading.Barrier(n_threads)
+
+        self._thread_ids = []
+        self._values = {}
+        self._entry_lock = threading.Lock()
+        self._entry_counter = 0
+
+    def __call__(self, **values):
+        """
+        Compares a set of values across threads.
+
+        For each value, type equality as well as comparison takes place. If any
+        of the values is a function, then address comparison is performed.
+        Also, if any of the values is a `numpy.ndarray`, then approximate
+        numerical comparison is performed.
+        """
+        tid = id(threading.current_thread())
+        self._entry_barrier.wait()
+        with self._entry_lock:
+            if self._entry_counter == 0:
+                # Reset state before comparison
+                self._barrier.reset()
+                self._reset_evt.clear()
+                self._thread_ids = []
+                self._values = {}
+                self._entry_barrier.reset()
+            self._entry_counter += 1
+
+        self._values[tid] = values
+        self._thread_ids.append(tid)
+        self._barrier.wait()
+
+        if tid == self._thread_ids[0]:
+            thread_ids = list(self._values)
+            try:
+                for value_name in values:
+                    for i in range(1, len(thread_ids)):
+                        tid_a = thread_ids[i - 1]
+                        tid_b = thread_ids[i]
+                        value_a = self._values[tid_a][value_name]
+                        value_b = self._values[tid_b][value_name]
+                        assert type(value_a) is type(value_b)
+                        if numpy_available and isinstance(value_a, np.ndarray):
+                            if len(value_a.shape) == 0:
+                                assert value_a == value_b
+                            else:
+                                assert np.allclose(value_a, value_b, equal_nan=True)
+                        elif isinstance(value_a, types.FunctionType):
+                            assert id(value_a) == id(value_b)
+                        elif value_a != value_a:
+                            assert value_b != value_b
+                        else:
+                            assert value_a == value_b
+            finally:
+                self._entry_counter = 0
+                self._reset_evt.set()
+        else:
+            self._reset_evt.wait()