Add back addBatchSystemFactory function (#3754)

src/toil/batchSystems/registry.py

@@ -13,6 +13,12 @@
 # limitations under the License.
 
 
+from typing import Callable, Type, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from toil.batchSystems.abstractBatchSystem import AbstractBatchSystem
+
+
 def gridengine_batch_system_factory():
     from toil.batchSystems.gridengine import GridEngineBatchSystem
     return GridEngineBatchSystem
@@ -71,3 +77,10 @@ def kubernetes_batch_system_factory():
 }
 BATCH_SYSTEMS = list(BATCH_SYSTEM_FACTORY_REGISTRY.keys())
 DEFAULT_BATCH_SYSTEM = 'single_machine'
+
+def addBatchSystemFactory(key: str, batchSystemFactory: Callable[[], Type['AbstractBatchSystem']]):
+    """
+    Adds a batch system to the registry for workflow-supplied batch systems.
+    """
+    BATCH_SYSTEMS.append(key)
+    BATCH_SYSTEM_FACTORY_REGISTRY[key] = batchSystemFactory

src/toil/test/batchSystems/batchSystemTest.py

@@ -34,6 +34,10 @@
 # protected by annotations.
 from toil.batchSystems.mesos.test import MesosTestSupport
 from toil.batchSystems.parasol import ParasolBatchSystem
+from toil.batchSystems.registry import (BATCH_SYSTEM_FACTORY_REGISTRY,
+                                        BATCH_SYSTEMS,
+                                        single_machine_batch_system_factory,
+                                        addBatchSystemFactory)
 from toil.test.batchSystems.parasolTestSupport import ParasolTestSupport
 from toil.batchSystems.singleMachine import SingleMachineBatchSystem
 from toil.common import Config, Toil
@@ -307,6 +311,14 @@ def _waitForJobsToStart(self, numJobs, tries=20):
                 time.sleep(1)
             return runningIDs
 
+        def testAddBatchSystemFactory(self):
+            def test_batch_system_factory():
+                return SingleMachineBatchSystem
+
+            addBatchSystemFactory('testBatchSystem', test_batch_system_factory)
+            assert ('testBatchSystem', test_batch_system_factory) in BATCH_SYSTEM_FACTORY_REGISTRY.items()
+            assert 'testBatchSystem' in BATCH_SYSTEMS
+
     class AbstractBatchSystemJobTest(ToilTest, metaclass=ABCMeta):
         """
         An abstract base class for batch system tests that use a full Toil workflow rather