Allow customizing env info

README.rst

@@ -166,9 +166,10 @@ These are the flags you can use when calling ``pytest`` in the command line:
   instead of saving them directly to the ``generated_image_dir``. Without this option,
   generated images are saved as ``generated_image_dir/<test_name>.png``; with this
   option enabled, they are instead saved as
-  ``<generated_image_dir>/<test_name>/<image_name>.png``, where the image name has
-  the format ``<system>_<python-version>_<pyvista-version>_<vtk-version>``. This can
-  be useful for providing context about how an image was generated.
+  ``<generated_image_dir>/<test_name>/<image_name>.png``, where the image name has the format
+  ``<os-version>_<machine>_<gpu-vendor>_<python-version>_<pyvista-version>_<vtk-version>_<using-ci>``.
+  This can be useful for providing context about how an image was generated. See the
+  ``Test specific flags`` section for customizing the info.
 
 * ``--failed_image_dir <DIR>`` dumps copies of cached and generated test images when
   there is a warning or error raised. This directory is useful for reviewing test
@@ -220,6 +221,29 @@ in the beginning of your test function.
   ``verify_image_cache`` fixture is used but no images are generated. The value of this
   flag takes precedence over the global flag by the same name (see above).
 
+* ``env_info``: Dataclass for controlling the environment info used to name the generated
+  test image(s) when the ``--generate_dirs`` option is used. The info can be test-specific
+  or can be modified globally by wrapping the ``verify_image_cache`` fixture, e.g.:
+
+  .. code-block:: python
+
+    @pytest.fixture(autouse=True)
+    def wrapped_verify_image_cache(verify_image_cache):
+        info = verify_image_cache.env_info
+
+        # NOTE: Default values are shown
+        info.prefix: str = ""  # Add a custom prefix
+        info.os: bool = True  # Show/hide the os version (e.g. ubuntu, macOS, Windows)
+        info.machine: bool = True  # Show/hide the machine info (e.g. arm64)
+        info.gpu: bool = True  # Show/hide the gpu vendor (e.g. NVIDIA)
+        info.python: bool = True  # Show/hide the python version
+        info.pyvista: bool = True  # Show/hide the pyvista version
+        info.vtk: bool = True  # Show/hide the vtk version
+        info.ci: bool = True  # Show/hide if generated in CI
+        info.suffix: str = ""  # Add a custom suffix
+
+        return verify_image_cache
+
 Documentation image tests
 -------------------------
 Unlike the unit tests, which use the ``verify_image_cache`` fixture to evaluate test

pytest_pyvista/pytest_pyvista.py

@@ -3,11 +3,13 @@
 from __future__ import annotations
 
 import contextlib
+from dataclasses import dataclass
 import importlib
 import json
 import os
 from pathlib import Path
 import platform
+import re
 import shutil
 import sys
 from typing import TYPE_CHECKING
@@ -28,23 +30,90 @@
 VISITED_CACHED_IMAGE_NAMES: set[str] = set()
 SKIPPED_CACHED_IMAGE_NAMES: set[str] = set()
 
-
-def _get_env_info() -> str:
-    system = platform.system()
-    if system == "Darwin":
-        system = "macOS"
-
-    return "_".join(
-        [
-            f"{system}-{platform.release()}",
-            f"py-{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}",
-            f"pyvista-{pyvista.__version__}",
-            f"vtk-{vtkmodules.__version__}",
+_GPU_VENDOR: list[str] = [""]  # Use a list so we can mutate the string globally
+
+
+@dataclass
+class _EnvInfo:
+    prefix: str = ""
+    os: bool = True
+    machine: bool = True
+    python: bool = True
+    pyvista: bool = True
+    vtk: bool = True
+    gpu: bool = True
+    ci: bool = True
+    suffix: str = ""
+
+    def __repr__(self) -> str:
+        os_info = _EnvInfo._get_os()
+        os_version = f"{os_info[0]}-{os_info[1]}" if self.os else ""
+        machine = f"{platform.machine()}" if self.machine else ""
+        gpu = f"gpu-{_EnvInfo._gpu_vendor()}" if self.gpu else ""
+        python_version = f"py-{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}" if self.python else ""
+        pyvista_version = f"pyvista-{pyvista.__version__}" if self.pyvista else ""
+        vtk_version = f"vtk-{vtkmodules.__version__}" if self.vtk else ""
+        ci = f"{'' if os.environ.get('CI', None) else 'no'}-CI" if self.ci else ""
+
+        values = [
+            f"{self.prefix}",
+            f"{os_version}",
+            f"{machine}",
+            f"{gpu}",
+            f"{python_version}",
+            f"{pyvista_version}",
+            f"{vtk_version}",
+            f"{ci}",
+            f"{self.suffix}",
         ]
-    )
+        return "_".join(val for val in values if val)
 
+    @staticmethod
+    def _get_os() -> tuple[str, str]:
+        system = platform.system()
+        if system == "Linux":
+            try:
+                name = platform.freedesktop_os_release()["ID"]
+                version = platform.freedesktop_os_release()["VERSION_ID"]
+            except AttributeError:
+                name = system
+                version = platform.release()
+            return name, version
+        name = "macOS" if system == "Darwin" else system
+        return name, platform.release()
 
-ENV_INFO = _get_env_info()
+    @staticmethod
+    def _gpu_vendor() -> str:
+        # Get cached value
+        if _GPU_VENDOR[0]:
+            return _GPU_VENDOR[0]
+
+        try:
+            vendor = pyvista.GPUInfo().vendor
+        except Exception:  # noqa: BLE001  # pragma: no cover
+            vendor = "UNKNOWN"
+
+        # Try to shorten vendor string
+        lower = vendor.lower()
+        if lower.startswith(nv := "nvidia"):  # pragma: no cover
+            text = nv
+        elif lower.startswith(amd := "amd"):  # pragma: no cover
+            text = amd
+        elif lower.startswith(ati := "ati"):  # pragma: no cover
+            text = ati
+        elif lower.startswith(mesa := "mesa"):
+            text = mesa
+        else:
+            text = vendor  # pragma: no cover
+        # Shorten original string and remove whitespace
+        vendor = vendor[: len(text)].replace(" ", "")
+        # Remove all potentially invalid/undesired filename characters
+        disallowed = r'[\\/:*?"<>|\s.\x00]'
+        vendor = re.sub(disallowed, "", vendor)
+
+        # Cache the value globally
+        _GPU_VENDOR[0] = vendor
+        return _GPU_VENDOR[0]
 
 
 class RegressionError(RuntimeError):
@@ -83,7 +152,7 @@ def pytest_addoption(parser) -> None:  # noqa: ANN001
     group.addoption(
         "--generate_subdirs",
         action="store_true",
-        help="Save generated images to sub-directories.",
+        help="Save generated images to sub-directories. The image names are determined by the environment info.",
     )
     group.addoption(
         "--add_missing_images",
@@ -223,6 +292,7 @@ class VerifyImageCache:
     add_missing_images = False
     reset_only_failed = False
     generate_subdirs = None
+    env_info: _EnvInfo
 
     def __init__(  # noqa: PLR0913
         self,
@@ -383,7 +453,9 @@ def remove_plotter_close_callback() -> None:
 
     def _save_generated_image(self, plotter: pyvista.Plotter, image_name: str, parent_dir: Path | None = None) -> None:
         parent = cast("Path", self.generated_image_dir) if parent_dir is None else parent_dir
-        generated_image_path = parent / Path(image_name).with_suffix("") / (ENV_INFO + ".png") if self.generate_subdirs else parent / image_name
+        generated_image_path = (
+            parent / Path(image_name).with_suffix("") / (str(self.env_info) + ".png") if self.generate_subdirs else parent / image_name
+        )
         generated_image_path.parent.mkdir(exist_ok=True, parents=True)
         plotter.screenshot(generated_image_path)
 
@@ -598,6 +670,7 @@ def verify_image_cache(
     VerifyImageCache.add_missing_images = pytestconfig.getoption("add_missing_images")
     VerifyImageCache.reset_only_failed = pytestconfig.getoption("reset_only_failed")
     VerifyImageCache.generate_subdirs = pytestconfig.getoption("generate_subdirs")
+    VerifyImageCache.env_info = _EnvInfo()
 
     cache_dir = cast("Path", _get_option_from_config_or_ini(pytestconfig, "image_cache_dir", is_dir=True))
     gen_dir = _get_option_from_config_or_ini(pytestconfig, "generated_image_dir", is_dir=True)

tests/test_pyvista.py

@@ -4,8 +4,10 @@
 
 from enum import Enum
 import filecmp
+import os
 from pathlib import Path
 import platform
+import re
 import shutil
 import sys
 from unittest import mock
@@ -14,7 +16,7 @@
 import pytest
 import pyvista as pv
 
-from pytest_pyvista.pytest_pyvista import _get_env_info
+from pytest_pyvista.pytest_pyvista import _EnvInfo
 
 pv.OFF_SCREEN = True
 
@@ -326,6 +328,9 @@ def test_generated_image_dir_commandline(pytester: pytest.Pytester, generate_sub
         import pyvista as pv
         pv.OFF_SCREEN = True
         def test_imcache(verify_image_cache):
+            verify_image_cache.env_info.prefix = 'prefix'
+            verify_image_cache.env_info.vtk = False
+            verify_image_cache.env_info.suffix = 'suffix'
             sphere = pv.Sphere()
             plotter = pv.Plotter()
             plotter.add_mesh(sphere, color="red")
@@ -339,8 +344,15 @@ def test_imcache(verify_image_cache):
     result = pytester.runpytest(*args)
     assert (pytester.path / "gen_dir").is_dir()
     if generate_subdirs:
-        with_suffix = _get_env_info() + ".png"
-        assert (pytester.path / "gen_dir" / "imcache" / with_suffix).is_file()
+        subdir = pytester.path / "gen_dir" / "imcache"
+        assert subdir.is_dir()
+        paths = list(subdir.iterdir())
+        assert len(paths) == 1
+        image_path = paths[0]
+        assert image_path.suffix == ".png"
+        assert image_path.name.startswith("prefix")
+        assert image_path.stem.endswith("suffix")
+        assert "vtk" not in image_path.name
     else:
         assert (pytester.path / "gen_dir" / "imcache.png").is_file()
     result.assert_outcomes(passed=1)
@@ -1013,3 +1025,68 @@ def test_imcache(verify_image_cache):
         assert from_test.is_file() == failed_image_dir
         if failed_image_dir:
             assert file_has_changed(str(from_test), str(from_cache))
+
+
+def test_env_info() -> None:
+    """Test env info dataclass."""
+    info = str(_EnvInfo())
+    os_info = _EnvInfo._get_os()  # noqa:SLF001
+    assert " " not in info
+    assert info.startswith(os_info[0] + "-" + os_info[1])
+    if platform.system() == "Linux" and sys.version_info >= (3, 10):
+        assert info.startswith("ubuntu")
+
+    # Generic regex for "_package-#.#.#" with optional suffix (like .dev0, .post1, etc.)
+    pattern = r"_[a-zA-Z]+-\d+\.\d+\.\d+(?:[a-zA-Z0-9\.]*)?"
+    matches = re.findall(pattern, info)
+
+    assert any(m.startswith("_py-") for m in matches), f"No pyvista version found in {info}"
+    assert any(m.startswith("_pyvista-") for m in matches), f"No pyvista version found in {info}"
+    assert any(m.startswith("_vtk-") for m in matches), f"No vtk version found in {info}"
+
+    assert any(f"gpu-{vendor.lower()}" in info.lower() for vendor in ["Apple", "NVIDIA", "Mesa", "AMD", "ATI"])
+
+    if os.environ.get("CI", None):
+        assert "no-CI" not in info
+        assert "CI" in info
+    else:
+        assert "no-CI" in info
+
+
+@pytest.mark.parametrize(
+    ("name", "value"),
+    [
+        ("os", _EnvInfo._get_os()[0]),  # noqa: SLF001
+        ("machine", platform.machine()),
+        ("gpu", "gpu-"),
+        ("python", "py-"),
+        ("pyvista", "pyvista-"),
+        ("vtk", "vtk-"),
+        ("ci", "-CI"),
+    ],
+)
+def test_env_info_exclude(name: str, value: str) -> None:
+    """Test removing parts of the env info."""
+    # Sanity check to ensure the value is there ordinarily
+    info = str(_EnvInfo())
+    assert value in info
+
+    # Test it's excluded
+    info = str(_EnvInfo(**{name: False}))
+    assert value not in info
+    assert "__" not in info
+
+
+def test_env_info_prefix_suffix() -> None:
+    """Test env info dataclass prefix and suffix."""
+    text = "foobar"
+    default = str(_EnvInfo())
+    sep = "_"
+    assert not default.startswith(sep)
+    assert not default.endswith(sep)
+
+    with_prefix = str(_EnvInfo(prefix=text))
+    assert f"{text}{sep}{default}" == with_prefix
+
+    with_suffix = str(_EnvInfo(suffix=text))
+    assert f"{default}{sep}{text}" == with_suffix