Allow custom scheduler names in --dist command line argument

changelog/970.feature

@@ -0,0 +1 @@
+`--dist` option allows custom scheduler names now.

src/xdist/dsession.py

@@ -14,13 +14,7 @@
 
 from xdist.remote import Producer
 from xdist.remote import WorkerInfo
-from xdist.scheduler import EachScheduling
-from xdist.scheduler import LoadFileScheduling
-from xdist.scheduler import LoadGroupScheduling
-from xdist.scheduler import LoadScheduling
-from xdist.scheduler import LoadScopeScheduling
 from xdist.scheduler import Scheduling
-from xdist.scheduler import WorkStealingScheduling
 from xdist.workermanage import NodeManager
 from xdist.workermanage import WorkerController
 
@@ -81,11 +75,18 @@ def report_line(self, line: str) -> None:
 
     @pytest.hookimpl(trylast=True)
     def pytest_sessionstart(self, session: pytest.Session) -> None:
-        """Creates and starts the nodes.
+        """Initializes the scheduler, creates and starts the nodes.
 
         The nodes are setup to put their events onto self.queue.  As
         soon as nodes start they will emit the worker_workerready event.
         """
+        self.sched = self.config.hook.pytest_xdist_make_scheduler(
+            config=self.config, log=self.log
+        )
+        if self.sched is None:
+            dist = self.config.getoption("dist")
+            raise pytest.UsageError(f"pytest-xdist: scheduler {dist!r} not found")
+
         self.nodemanager = NodeManager(self.config)
         nodes = self.nodemanager.setup_nodes(putevent=self.queue.put)
         self._active_nodes.update(nodes)
@@ -104,34 +105,8 @@ def pytest_collection(self) -> bool:
         # prohibit collection of test items in controller process
         return True
 
-    @pytest.hookimpl(trylast=True)
-    def pytest_xdist_make_scheduler(
-        self,
-        config: pytest.Config,
-        log: Producer,
-    ) -> Scheduling | None:
-        dist = config.getvalue("dist")
-        if dist == "each":
-            return EachScheduling(config, log)
-        if dist == "load":
-            return LoadScheduling(config, log)
-        if dist == "loadscope":
-            return LoadScopeScheduling(config, log)
-        if dist == "loadfile":
-            return LoadFileScheduling(config, log)
-        if dist == "loadgroup":
-            return LoadGroupScheduling(config, log)
-        if dist == "worksteal":
-            return WorkStealingScheduling(config, log)
-        return None
-
     @pytest.hookimpl
     def pytest_runtestloop(self) -> bool:
-        self.sched = self.config.hook.pytest_xdist_make_scheduler(
-            config=self.config, log=self.log
-        )
-        assert self.sched is not None
-
         self.shouldstop = False
         pending_exception = None
         while not self.session_finished:

src/xdist/plugin.py

@@ -101,15 +101,6 @@ def pytest_addoption(parser: pytest.Parser) -> None:
         "--dist",
         metavar="distmode",
         action="store",
-        choices=[
-            "each",
-            "load",
-            "loadscope",
-            "loadfile",
-            "loadgroup",
-            "worksteal",
-            "no",
-        ],
         dest="dist",
         default="no",
         help=(
@@ -235,6 +226,13 @@ def pytest_configure(config: pytest.Config) -> None:
     # Create the distributed session in case we have a valid distribution
     # mode and test environments.
     if _is_distribution_mode(config):
+        config.pluginmanager.import_plugin("xdist.scheduler.each")
+        config.pluginmanager.import_plugin("xdist.scheduler.load")
+        config.pluginmanager.import_plugin("xdist.scheduler.loadfile")
+        config.pluginmanager.import_plugin("xdist.scheduler.loadgroup")
+        config.pluginmanager.import_plugin("xdist.scheduler.loadscope")
+        config.pluginmanager.import_plugin("xdist.scheduler.worksteal")
+
         from xdist.dsession import DSession
 
         session = DSession(config)

src/xdist/scheduler/each.py

@@ -6,6 +6,7 @@
 
 from xdist.remote import Producer
 from xdist.report import report_collection_diff
+from xdist.scheduler.protocol import Scheduling
 from xdist.workermanage import parse_spec_config
 from xdist.workermanage import WorkerController
 
@@ -150,3 +151,14 @@ def schedule(self) -> None:
             else:
                 node.send_runtest_some(pending)
             self._started.append(node)
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_xdist_make_scheduler(
+    config: pytest.Config,
+    log: Producer,
+) -> Scheduling | None:
+    if config.getoption("dist") == "each":
+        return EachScheduling(config, log)
+    else:
+        return None

src/xdist/scheduler/load.py

@@ -7,6 +7,7 @@
 
 from xdist.remote import Producer
 from xdist.report import report_collection_diff
+from xdist.scheduler.protocol import Scheduling
 from xdist.workermanage import parse_spec_config
 from xdist.workermanage import WorkerController
 
@@ -333,3 +334,14 @@ def _check_nodes_have_same_collection(self) -> bool:
                     self.config.hook.pytest_collectreport(report=rep)
 
         return same_collection
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_xdist_make_scheduler(
+    config: pytest.Config,
+    log: Producer,
+) -> Scheduling | None:
+    if config.getoption("dist") == "load":
+        return LoadScheduling(config, log)
+    else:
+        return None

src/xdist/scheduler/loadfile.py

@@ -3,6 +3,7 @@
 import pytest
 
 from xdist.remote import Producer
+from xdist.scheduler.protocol import Scheduling
 
 from .loadscope import LoadScopeScheduling
 
@@ -58,3 +59,14 @@ def _split_scope(self, nodeid: str) -> str:
             example/loadsuite/epsilon/__init__.py
         """
         return nodeid.split("::", 1)[0]
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_xdist_make_scheduler(
+    config: pytest.Config,
+    log: Producer,
+) -> Scheduling | None:
+    if config.getoption("dist") == "loadfile":
+        return LoadFileScheduling(config, log)
+    else:
+        return None

src/xdist/scheduler/loadgroup.py

@@ -3,6 +3,7 @@
 import pytest
 
 from xdist.remote import Producer
+from xdist.scheduler.protocol import Scheduling
 
 from .loadscope import LoadScopeScheduling
 
@@ -57,3 +58,14 @@ def _split_scope(self, nodeid: str) -> str:
             return nodeid.split("@")[-1]
         else:
             return nodeid
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_xdist_make_scheduler(
+    config: pytest.Config,
+    log: Producer,
+) -> Scheduling | None:
+    if config.getoption("dist") == "loadgroup":
+        return LoadGroupScheduling(config, log)
+    else:
+        return None

src/xdist/scheduler/loadscope.py

@@ -8,6 +8,7 @@
 
 from xdist.remote import Producer
 from xdist.report import report_collection_diff
+from xdist.scheduler.protocol import Scheduling
 from xdist.workermanage import parse_spec_config
 from xdist.workermanage import WorkerController
 
@@ -432,3 +433,14 @@ def _check_nodes_have_same_collection(self) -> bool:
             self.config.hook.pytest_collectreport(report=rep)
 
         return same_collection
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_xdist_make_scheduler(
+    config: pytest.Config,
+    log: Producer,
+) -> Scheduling | None:
+    if config.getoption("dist") == "loadscope":
+        return LoadScopeScheduling(config, log)
+    else:
+        return None

src/xdist/scheduler/worksteal.py

@@ -7,6 +7,7 @@
 
 from xdist.remote import Producer
 from xdist.report import report_collection_diff
+from xdist.scheduler.protocol import Scheduling
 from xdist.workermanage import parse_spec_config
 from xdist.workermanage import WorkerController
 
@@ -343,3 +344,14 @@ def _check_nodes_have_same_collection(self) -> bool:
                     self.config.hook.pytest_collectreport(report=rep)
 
         return same_collection
+
+
+@pytest.hookimpl(trylast=True)
+def pytest_xdist_make_scheduler(
+    config: pytest.Config,
+    log: Producer,
+) -> Scheduling | None:
+    if config.getoption("dist") == "worksteal":
+        return WorkStealingScheduling(config, log)
+    else:
+        return None

testing/acceptance_test.py

@@ -1641,3 +1641,15 @@ def test():
     )
     result = pytester.runpytest()
     assert result.ret == 0
+
+
+def test_dist_validation(pytester: pytest.Pytester) -> None:
+    """Should exit early if incorrect --dist value is specified."""
+    f = pytester.makepyfile(
+        """
+        assert 0
+        """
+    )
+    result = pytester.runpytest(f, "-n1", "--dist=invalid")
+    assert result.ret == pytest.ExitCode.USAGE_ERROR
+    result.stderr.fnmatch_lines(["ERROR: pytest-xdist: scheduler 'invalid' not found"])