Get -nauto default from PYTEST_XDIST_AUTO_NUM_WORKERS (#829)

* Get `-nauto` default from `PYTEST_XDIST_AUTO_NUM_WORKERS`

A few additional tests for existing functionality are added too.
And the ``-n logical`` option is documented.

Fixes: #792

Author: encukou

changelog/792.feature

@@ -0,0 +1,2 @@
+The environment variable ``PYTEST_XDIST_AUTO_NUM_WORKERS`` can now be used to
+specify the default for ``-n auto`` and ``-n logical``.

changelog/829.doc

@@ -0,0 +1 @@
+Document the ``-n logical`` option.

docs/distribution.rst

@@ -12,13 +12,29 @@ noticeable amount of time.
 
 With ``-n auto``, pytest-xdist will use as many processes as your computer
 has CPU cores.
+
+Use ``-n logical`` to use the number of *logical* CPU cores rather than
+physical ones. This currently requires the ``psutils`` package to be installed;
+if it is not, pytest-xdist will fall back to ``-n auto`` behavior.
+
 Pass a number, e.g. ``-n 8``, to specify the number of processes explicitly.
 
-To specify a different meaning for ``-n auto`` for your tests,
-you can implement the ``pytest_xdist_auto_num_workers``
-`pytest hook <https://docs.pytest.org/en/latest/how-to/writing_plugins.html>`__
-(a function named ``pytest_xdist_auto_num_workers`` in e.g. ``conftest.py``)
-that returns the number of processes to use.
+To specify a different meaning for ``-n auto`` and ``-n logical`` for your
+tests, you can:
+
+* Set the environment variable ``PYTEST_XDIST_AUTO_NUM_WORKERS`` to the
+  desired number of processes.
+
+* Implement the ``pytest_xdist_auto_num_workers``
+  `pytest hook <https://docs.pytest.org/en/latest/how-to/writing_plugins.html>`__
+  (a ``pytest_xdist_auto_num_workers(config)`` function in e.g. ``conftest.py``)
+  that returns the number of processes to use.
+  The hook can use ``config.option.numprocesses`` to determine if the user
+  asked for ``"auto"`` or ``"logical"``, and it can return ``None`` to fall
+  back to the default.
+
+If both the hook and environment variable are specified, the hook takes
+priority.
 
 
 Parallelization can be configured further with these options:

src/xdist/plugin.py

@@ -1,6 +1,7 @@
 import os
 import uuid
 import sys
+import warnings
 
 import pytest
 
@@ -12,6 +13,13 @@
 
 @pytest.hookimpl
 def pytest_xdist_auto_num_workers(config):
+    env_var = os.environ.get("PYTEST_XDIST_AUTO_NUM_WORKERS")
+    if env_var:
+        try:
+            return int(env_var)
+        except ValueError:
+            warnings.warn("PYTEST_XDIST_AUTO_NUM_WORKERS is not a number: {env_var!r}. Ignoring it.")
+
     try:
         import psutil
     except ImportError:

testing/test_plugin.py

@@ -1,12 +1,22 @@
 from contextlib import suppress
 from pathlib import Path
+import sys
+import os
 
 import execnet
 from xdist.workermanage import NodeManager
 
 import pytest
 
 
+@pytest.fixture
+def monkeypatch_3_cpus(monkeypatch: pytest.MonkeyPatch):
+    """Make pytest-xdist believe the system has 3 CPUs"""
+    monkeypatch.setitem(sys.modules, "psutil", None)  # block import
+    monkeypatch.delattr(os, "sched_getaffinity", raising=False)
+    monkeypatch.setattr(os, "cpu_count", lambda: 3)
+
+
 def test_dist_incompatibility_messages(pytester: pytest.Pytester) -> None:
     result = pytester.runpytest("--pdb", "--looponfail")
     assert result.ret != 0
@@ -41,7 +51,6 @@ def test_dist_options(pytester: pytest.Pytester) -> None:
 def test_auto_detect_cpus(
     pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch
 ) -> None:
-    import os
     from xdist.plugin import pytest_cmdline_main as check_options
 
     with suppress(ImportError):
@@ -102,6 +111,20 @@ def test_auto_detect_cpus_psutil(
     assert config.getoption("numprocesses") == 84
 
 
+def test_auto_detect_cpus_os(
+    pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch, monkeypatch_3_cpus
+) -> None:
+    from xdist.plugin import pytest_cmdline_main as check_options
+
+    config = pytester.parseconfigure("-nauto")
+    check_options(config)
+    assert config.getoption("numprocesses") == 3
+
+    config = pytester.parseconfigure("-nlogical")
+    check_options(config)
+    assert config.getoption("numprocesses") == 3
+
+
 def test_hook_auto_num_workers(
     pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch
 ) -> None:
@@ -122,6 +145,105 @@ def pytest_xdist_auto_num_workers():
     assert config.getoption("numprocesses") == 42
 
 
+def test_hook_auto_num_workers_arg(
+    pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    # config.option.numprocesses is a pytest feature,
+    # but we document it so let's test it.
+    from xdist.plugin import pytest_cmdline_main as check_options
+
+    pytester.makeconftest(
+        """
+        def pytest_xdist_auto_num_workers(config):
+            if config.option.numprocesses == 'auto':
+                return 42
+            if config.option.numprocesses == 'logical':
+                return 8
+    """
+    )
+    config = pytester.parseconfigure("-nauto")
+    check_options(config)
+    assert config.getoption("numprocesses") == 42
+
+    config = pytester.parseconfigure("-nlogical")
+    check_options(config)
+    assert config.getoption("numprocesses") == 8
+
+
+def test_hook_auto_num_workers_none(
+    pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch, monkeypatch_3_cpus
+) -> None:
+    # Returning None from a hook to skip it is pytest behavior,
+    # but we document it so let's test it.
+    from xdist.plugin import pytest_cmdline_main as check_options
+
+    pytester.makeconftest(
+        """
+        def pytest_xdist_auto_num_workers():
+            return None
+    """
+    )
+    config = pytester.parseconfigure("-nauto")
+    check_options(config)
+    assert config.getoption("numprocesses") == 3
+
+    monkeypatch.setenv("PYTEST_XDIST_AUTO_NUM_WORKERS", "5")
+
+    config = pytester.parseconfigure("-nauto")
+    check_options(config)
+    assert config.getoption("numprocesses") == 5
+
+
+def test_envvar_auto_num_workers(
+    pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    from xdist.plugin import pytest_cmdline_main as check_options
+
+    monkeypatch.setenv("PYTEST_XDIST_AUTO_NUM_WORKERS", "7")
+
+    config = pytester.parseconfigure("-nauto")
+    check_options(config)
+    assert config.getoption("numprocesses") == 7
+
+    config = pytester.parseconfigure("-nlogical")
+    check_options(config)
+    assert config.getoption("numprocesses") == 7
+
+
+def test_envvar_auto_num_workers_warn(
+    pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch, monkeypatch_3_cpus
+) -> None:
+    from xdist.plugin import pytest_cmdline_main as check_options
+
+    monkeypatch.setenv("PYTEST_XDIST_AUTO_NUM_WORKERS", "fourscore")
+
+    config = pytester.parseconfigure("-nauto")
+    with pytest.warns(UserWarning):
+        check_options(config)
+    assert config.getoption("numprocesses") == 3
+
+
+def test_auto_num_workers_hook_overrides_envvar(
+    pytester: pytest.Pytester, monkeypatch: pytest.MonkeyPatch, monkeypatch_3_cpus
+) -> None:
+    from xdist.plugin import pytest_cmdline_main as check_options
+
+    monkeypatch.setenv("PYTEST_XDIST_AUTO_NUM_WORKERS", "987")
+    pytester.makeconftest(
+        """
+        def pytest_xdist_auto_num_workers():
+            return 2
+    """
+    )
+    config = pytester.parseconfigure("-nauto")
+    check_options(config)
+    assert config.getoption("numprocesses") == 2
+
+    config = pytester.parseconfigure("-nauto")
+    check_options(config)
+    assert config.getoption("numprocesses") == 2
+
+
 def test_dsession_with_collect_only(pytester: pytest.Pytester) -> None:
     from xdist.plugin import pytest_cmdline_main as check_options
     from xdist.plugin import pytest_configure as configure