Refactor plugin logic with a cleaner structure

Author: art049

.github/workflows/ci.yml

@@ -39,6 +39,5 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .[dev]
-          pip uninstall pytest-benchmark -y
       - name: Run tests
         run: pytest

pyproject.toml

@@ -44,7 +44,6 @@ dev = [
     "pytest ~= 7.0",
     "pytest-cov ~= 4.0.0",
     "ruff ~= 0.0.100",
-    "pytest-benchmark>=3.4.1",
 ]
 
 [project.entry-points]

src/pytest_codspeed/_wrapper/__init__.py

@@ -4,7 +4,7 @@
 from cffi import FFI  # type: ignore
 
 if TYPE_CHECKING:
-    from .wrapper import lib as _lib
+    from .wrapper import lib as LibType
 
 _wrapper_dir = os.path.dirname(os.path.abspath(__file__))
 
@@ -21,7 +21,7 @@ def _get_ffi():
     return ffi
 
 
-def get_lib() -> "_lib":
+def get_lib() -> "LibType":
     try:
         ffi = _get_ffi()
         ffi.compile(

src/pytest_codspeed/_wrapper/wrapper.pyi

@@ -1,7 +1,13 @@
 class lib:
-    def start_instrumentation(self) -> None: ...
-    def stop_instrumentation(self) -> None: ...
-    def dump_stats(self) -> None: ...
-    def dump_stats_at(self, trigger: bytes) -> None: ...
-    def zero_stats(self) -> None: ...
-    def toggle_collect(self) -> None: ...
+    @staticmethod
+    def start_instrumentation() -> None: ...
+    @staticmethod
+    def stop_instrumentation() -> None: ...
+    @staticmethod
+    def dump_stats() -> None: ...
+    @staticmethod
+    def dump_stats_at(trigger: bytes) -> None: ...
+    @staticmethod
+    def zero_stats() -> None: ...
+    @staticmethod
+    def toggle_collect() -> None: ...

src/pytest_codspeed/plugin.py

@@ -1,19 +1,14 @@
 import os
-from typing import Any, Callable, List
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Callable, List, Optional
 
 import pytest
 
 from . import __version__
 from ._wrapper import get_lib
 
-lib = get_lib()
-
-_benchmark_count = 0
-
-
-@pytest.hookimpl(trylast=True)
-def pytest_report_header(config: "pytest.Config"):
-    return f"codspeed: {__version__}"
+if TYPE_CHECKING:
+    from ._wrapper import LibType
 
 
 @pytest.hookimpl(trylast=True)
@@ -27,6 +22,21 @@ def pytest_addoption(parser: "pytest.Parser"):
     )
 
 
+@dataclass(unsafe_hash=True)
+class CodSpeedPlugin:
+    is_codspeed_enabled: bool
+    should_measure: bool
+    lib: Optional["LibType"] = None
+    benchmark_count: int = 0
+
+
+PLUGIN_NAME = "codspeed_plugin"
+
+
+def get_plugin(config: "pytest.Config") -> "CodSpeedPlugin":
+    return config.pluginmanager.get_plugin(PLUGIN_NAME)
+
+
 @pytest.hookimpl()
 def pytest_configure(config: "pytest.Config"):
     config.addinivalue_line(
@@ -35,10 +45,26 @@ def pytest_configure(config: "pytest.Config"):
     config.addinivalue_line(
         "markers", "benchmark: mark an entire test for codspeed benchmarking"
     )
+    plugin = CodSpeedPlugin(
+        is_codspeed_enabled=config.getoption("--codspeed")
+        or os.environ.get("CODSPEED_ENV") is not None,
+        should_measure=os.environ.get("CODSPEED_ENV") is not None,
+    )
+    if plugin.should_measure:
+        plugin.lib = get_lib()
+    config.pluginmanager.register(plugin, PLUGIN_NAME)
 
 
-def is_benchmark_enabled(config: "pytest.Config") -> bool:
-    return config.getoption("--codspeed") or os.environ.get("CODSPEED_ENV") is not None
+@pytest.hookimpl(trylast=True)
+def pytest_report_header(config: "pytest.Config"):
+    out = [f"codspeed: {__version__}"]
+    plugin = get_plugin(config)
+    if plugin.is_codspeed_enabled and not plugin.should_measure:
+        out.append(
+            "NOTICE: codspeed is enabled, but no performance measurement"
+            " will be made since it's running in an unknown environment."
+        )
+    return "\n".join(out)
 
 
 def should_benchmark_item(item: "pytest.Item") -> bool:
@@ -51,16 +77,17 @@ def should_benchmark_item(item: "pytest.Item") -> bool:
 
 @pytest.hookimpl()
 def pytest_sessionstart(session: "pytest.Session"):
-    if is_benchmark_enabled(session.config):
-        global _benchmark_count
-        _benchmark_count = 0
+    plugin = get_plugin(session.config)
+    if plugin.is_codspeed_enabled:
+        plugin.benchmark_count = 0
 
 
 @pytest.hookimpl(trylast=True)
 def pytest_collection_modifyitems(
     session: "pytest.Session", config: "pytest.Config", items: "List[pytest.Item]"
 ):
-    if is_benchmark_enabled(session.config):
+    plugin = get_plugin(config)
+    if plugin.is_codspeed_enabled:
         deselected = []
         selected = []
         for item in items:
@@ -74,45 +101,51 @@ def pytest_collection_modifyitems(
 
 @pytest.hookimpl()
 def pytest_runtest_call(item: "pytest.Item"):
-    if not is_benchmark_enabled(item.config) or not should_benchmark_item(item):
+    plugin = get_plugin(item.config)
+    if not plugin.is_codspeed_enabled or not should_benchmark_item(item):
         item.runtest()
     else:
-        global _benchmark_count
-        _benchmark_count += 1
+        plugin.benchmark_count += 1
         if "benchmark" in getattr(item, "fixturenames", []):
+            # This is a benchmark fixture, so the measurement is done by the fixture
+            item.runtest()
+        elif not plugin.should_measure:
             item.runtest()
         else:
-            lib.zero_stats()
-            lib.start_instrumentation()
+            assert plugin.lib is not None
+            plugin.lib.zero_stats()
+            plugin.lib.start_instrumentation()
             item.runtest()
-            lib.stop_instrumentation()
-            lib.dump_stats_at(f"{item.nodeid}".encode("ascii"))
+            plugin.lib.stop_instrumentation()
+            plugin.lib.dump_stats_at(f"{item.nodeid}".encode("ascii"))
 
 
 @pytest.hookimpl()
 def pytest_sessionfinish(session: "pytest.Session", exitstatus):
-    if is_benchmark_enabled(session.config):
+    plugin = get_plugin(session.config)
+    if plugin.is_codspeed_enabled:
         reporter = session.config.pluginmanager.get_plugin("terminalreporter")
-        reporter.write_sep("=", f"{_benchmark_count} benchmarked")
-
-
-@pytest.fixture(scope="session")
-def _is_benchmark_enabled(request: "pytest.FixtureRequest") -> bool:
-    return is_benchmark_enabled(request.config)
+        count_suffix = "benchmarked" if plugin.should_measure else "benchmark tested"
+        reporter.write_sep(
+            "=",
+            f"{plugin.benchmark_count} {count_suffix}",
+        )
 
 
 @pytest.fixture
-def codspeed_benchmark(
-    request: "pytest.FixtureRequest", _is_benchmark_enabled: bool
-) -> Callable:
+def codspeed_benchmark(request: "pytest.FixtureRequest") -> Callable:
+    plugin = get_plugin(request.config)
+
     def run(func: Callable[..., Any], *args: Any):
-        if _is_benchmark_enabled:
-            lib.zero_stats()
-            lib.start_instrumentation()
-        func(*args)
-        if _is_benchmark_enabled:
-            lib.stop_instrumentation()
-            lib.dump_stats_at(f"{request.node.nodeid}".encode("ascii"))
+        if plugin.is_codspeed_enabled and plugin.should_measure:
+            assert plugin.lib is not None
+            plugin.lib.zero_stats()
+            plugin.lib.start_instrumentation()
+            func(*args)
+            plugin.lib.stop_instrumentation()
+            plugin.lib.dump_stats_at(f"{request.node.nodeid}".encode("ascii"))
+        else:
+            func(*args)
 
     return run
 

tests/test_pytest_plugin.py

@@ -1,9 +1,19 @@
 import os
+from contextlib import contextmanager
 
 import pytest
 
 
-def test_plugin_enabled_cli(pytester: pytest.Pytester) -> None:
+@contextmanager
+def codspeed_env():
+    os.environ["CODSPEED_ENV"] = "1"
+    try:
+        yield
+    finally:
+        del os.environ["CODSPEED_ENV"]
+
+
+def test_plugin_enabled_without_env(pytester: pytest.Pytester) -> None:
     pytester.makepyfile(
         """
         def test_some_addition_performance(benchmark):
@@ -13,10 +23,30 @@ def _():
         """
     )
     result = pytester.runpytest("--codspeed")
+    result.stdout.fnmatch_lines(
+        [
+            "NOTICE: codspeed is enabled, but no performance measurement will be made*",
+            "*1 benchmark tested*",
+            "*1 passed*",
+        ]
+    )
+
+
+def test_plugin_enabled_by_env(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile(
+        """
+        def test_some_addition_performance(benchmark):
+            @benchmark
+            def _():
+                return 1 + 1
+        """
+    )
+    with codspeed_env():
+        result = pytester.runpytest()
     result.stdout.fnmatch_lines(["*1 benchmarked*", "*1 passed*"])
 
 
-def test_plugin_enabled_env(pytester: pytest.Pytester) -> None:
+def test_plugin_enabled_and_env(pytester: pytest.Pytester) -> None:
     pytester.makepyfile(
         """
         def test_some_addition_performance(benchmark):
@@ -25,11 +55,9 @@ def _():
                 return 1 + 1
         """
     )
-    os.environ["CODSPEED_ENV"] = "1"
-    pytester.runpytest()
-    result = pytester.runpytest()
+    with codspeed_env():
+        result = pytester.runpytest("--codspeed")
     result.stdout.fnmatch_lines(["*1 benchmarked*", "*1 passed*"])
-    del os.environ["CODSPEED_ENV"]
 
 
 def test_plugin_disabled(pytester: pytest.Pytester) -> None:
@@ -52,7 +80,8 @@ def test_some_addition_performance():
             return 1 + 1
         """
     )
-    result = pytester.runpytest("--codspeed")
+    with codspeed_env():
+        result = pytester.runpytest("--codspeed")
     result.stdout.fnmatch_lines(["*0 benchmarked*", "*1 deselected*"])