Allow custom scheduler class

xdist/dsession.py

@@ -1,4 +1,5 @@
 import difflib
+import importlib
 import itertools
 from _pytest.runner import CollectReport
 
@@ -24,7 +25,7 @@ class EachScheduling:
     assigned the remaining items from the removed node.
     """
 
-    def __init__(self, numnodes, log=None):
+    def __init__(self, numnodes, log=None, config=None):
         self.numnodes = numnodes
         self.node2collection = {}
         self.node2pending = {}
@@ -485,6 +486,17 @@ class DSession:
             self.trdist = TerminalDistReporter(config)
             config.pluginmanager.register(self.trdist, "terminaldistreporter")
 
+        self.sched = None
+        dist = self.config.getvalue("dist")
+        if dist == "load":
+            self.ScheduleCls = LoadScheduling
+        elif dist == "each":
+            self.ScheduleCls = EachScheduling
+        else:
+            mod_name, cls_name = dist.rsplit('.', 1)
+            mod = importlib.import_module(mod_name)
+            self.ScheduleCls = getattr(mod, cls_name)
+
     @property
     def session_finished(self):
         """Return True if the distributed session has finished
@@ -523,15 +535,9 @@ class DSession:
 
     def pytest_runtestloop(self):
         numnodes = len(self.nodemanager.specs)
-        dist = self.config.getvalue("dist")
-        if dist == "load":
-            self.sched = LoadScheduling(numnodes, log=self.log,
-                                        config=self.config)
-        elif dist == "each":
-            self.sched = EachScheduling(numnodes, log=self.log)
-        else:
-            assert 0, dist
         self.shouldstop = False
+        self.sched = self.ScheduleCls(numnodes, log=self.log,
+                                      config=self.config)
         while not self.session_finished:
             self.loop_once()
             if self.shouldstop:

xdist/plugin.py

@@ -37,6 +37,8 @@ def pytest_addoption(parser):
               "each: send each test to each available environment.\n\n"
               "load: send each test to available environment.\n\n"
               "(default) no: run tests inprocess, don't distribute."))
+    group._addoption('--dc', dest='distcustom',
+                    help="Custom scheduler implementation")
     group._addoption(
         '--tx', dest="tx", action="append", default=[],
         metavar="xspec",
@@ -99,6 +101,8 @@ def pytest_cmdline_main(config):
         config.option.tx = ['popen'] * config.option.numprocesses
     if config.option.distload:
         config.option.dist = "load"
+    if config.option.distcustom:
+        config.option.dist = config.option.distcustom
     val = config.getvalue
     if not val("collectonly"):
         usepdb = config.getoption('usepdb')  # a core option

xdist/remote.py

@@ -131,6 +131,7 @@ def remote_initconfig(option_dict, args):
     config.option.usepdb = False
     config.option.dist = "no"
     config.option.distload = False
+    config.option.distcustom = None
     config.option.numprocesses = None
     config.args = args
     return config