toil 6.1.0a1__py3-none-any.whl → 8.0.0__py3-none-any.whl

toil/batchSystems/registry.py

@@ -16,7 +16,8 @@ import importlib
 import logging
 import pkgutil
 import warnings
-from typing import TYPE_CHECKING, Callable, Dict, List, Sequence, Tuple, Type
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Callable
 
 from toil.lib.compatibility import deprecated
 from toil.lib.memoize import memoize
@@ -30,7 +31,10 @@ logger = logging.getLogger(__name__)
 # Plugin system/API
 #####
 
-def add_batch_system_factory(key: str, class_factory: Callable[[], Type['AbstractBatchSystem']]):
+
+def add_batch_system_factory(
+    key: str, class_factory: Callable[[], type["AbstractBatchSystem"]]
+):
     """
     Adds a batch system to the registry for workflow or plugin-supplied batch systems.
 
@@ -39,6 +43,7 @@ def add_batch_system_factory(key: str, class_factory: Callable[[], Type['Abstrac
     _registry_keys.append(key)
     _registry[key] = class_factory
 
+
 def get_batch_systems() -> Sequence[str]:
     """
     Get the names of all the availsble batch systems.
@@ -47,7 +52,8 @@ def get_batch_systems() -> Sequence[str]:
 
     return _registry_keys
 
-def get_batch_system(key: str) -> Type['AbstractBatchSystem']:
+
+def get_batch_system(key: str) -> type["AbstractBatchSystem"]:
     """
     Get a batch system class by name.
 
@@ -58,68 +64,81 @@ def get_batch_system(key: str) -> Type['AbstractBatchSystem']:
     return _registry[key]()
 
 
-DEFAULT_BATCH_SYSTEM = 'single_machine'
+DEFAULT_BATCH_SYSTEM = "single_machine"
 
 #####
 # Built-in batch systems
 #####
 
+
 def aws_batch_batch_system_factory():
     from toil.batchSystems.awsBatch import AWSBatchBatchSystem
+
     return AWSBatchBatchSystem
 
+
 def gridengine_batch_system_factory():
     from toil.batchSystems.gridengine import GridEngineBatchSystem
+
     return GridEngineBatchSystem
 
 
 def lsf_batch_system_factory():
     from toil.batchSystems.lsf import LSFBatchSystem
+
     return LSFBatchSystem
 
 
 def single_machine_batch_system_factory():
     from toil.batchSystems.singleMachine import SingleMachineBatchSystem
+
     return SingleMachineBatchSystem
 
 
 def mesos_batch_system_factory():
     from toil.batchSystems.mesos.batchSystem import MesosBatchSystem
+
     return MesosBatchSystem
 
 
 def slurm_batch_system_factory():
     from toil.batchSystems.slurm import SlurmBatchSystem
+
     return SlurmBatchSystem
 
+
 def torque_batch_system_factory():
     from toil.batchSystems.torque import TorqueBatchSystem
+
     return TorqueBatchSystem
 
 
 def htcondor_batch_system_factory():
     from toil.batchSystems.htcondor import HTCondorBatchSystem
+
     return HTCondorBatchSystem
 
 
 def kubernetes_batch_system_factory():
     from toil.batchSystems.kubernetes import KubernetesBatchSystem
+
     return KubernetesBatchSystem
 
+
 #####
 # Registry implementation
 #####
 
-_registry: Dict[str, Callable[[], Type["AbstractBatchSystem"]]] = {
-    'aws_batch'      : aws_batch_batch_system_factory,
-    'single_machine' : single_machine_batch_system_factory,
-    'grid_engine'    : gridengine_batch_system_factory,
-    'lsf'            : lsf_batch_system_factory,
-    'mesos'          : mesos_batch_system_factory,
-    'slurm'          : slurm_batch_system_factory,
-    'torque'         : torque_batch_system_factory,
-    'htcondor'       : htcondor_batch_system_factory,
-    'kubernetes'     : kubernetes_batch_system_factory
+_registry: dict[str, Callable[[], type["AbstractBatchSystem"]]] = {
+    "aws_batch": aws_batch_batch_system_factory,
+    "single_machine": single_machine_batch_system_factory,
+    "grid_engine": gridengine_batch_system_factory,
+    "lsf": lsf_batch_system_factory,
+    "mesos": mesos_batch_system_factory,
+    "slurm": slurm_batch_system_factory,
+    "torque": torque_batch_system_factory,
+    "htcondor": htcondor_batch_system_factory,
+    "kubernetes": kubernetes_batch_system_factory,
 }
 _registry_keys = list(_registry.keys())
 
@@ -127,6 +146,7 @@ _registry_keys = list(_registry.keys())
 # add_batch_system_factory()
 _PLUGIN_NAME_PREFIX = "toil_batch_system_"
 
+
 @memoize
 def _load_all_plugins() -> None:
     """
@@ -139,6 +159,7 @@ def _load_all_plugins() -> None:
             # If it is a Toil batch system plugin, import it
             importlib.import_module(name)
 
+
 #####
 # Deprecated API
 #####
@@ -155,17 +176,24 @@ def __getattr__(name):
     See <https://stackoverflow.com/a/48242860>.
     """
     if name == "BATCH_SYSTEM_FACTORY_REGISTRY":
-        warnings.warn("BATCH_SYSTEM_FACTORY_REGISTRY is deprecated; use get_batch_system() or add_batch_system_factory()", DeprecationWarning)
+        warnings.warn(
+            "BATCH_SYSTEM_FACTORY_REGISTRY is deprecated; use get_batch_system() or add_batch_system_factory()",
+            DeprecationWarning,
+        )
         return _registry
     elif name == "BATCH_SYSTEMS":
-        warnings.warn("BATCH_SYSTEMS is deprecated; use get_batch_systems()", DeprecationWarning)
+        warnings.warn(
+            "BATCH_SYSTEMS is deprecated; use get_batch_systems()", DeprecationWarning
+        )
         return _registry_keys
     else:
         raise AttributeError(f"Module {__name__} ahs no attribute {name}")
 
 
 @deprecated(new_function_name="add_batch_system_factory")
-def addBatchSystemFactory(key: str, batchSystemFactory: Callable[[], Type['AbstractBatchSystem']]):
+def addBatchSystemFactory(
+    key: str, batchSystemFactory: Callable[[], type["AbstractBatchSystem"]]
+):
     """
     Deprecated method to add a batch system.
     """
@@ -180,7 +208,10 @@ def addBatchSystemFactory(key: str, batchSystemFactory: Callable[[], Type['Abstr
 # the globals because module-level globals are their own references, so we
 # can't touch this module's global name bindings from a client module.
 
-def save_batch_system_plugin_state() -> Tuple[List[str], Dict[str, Callable[[], Type['AbstractBatchSystem']]]]:
+
+def save_batch_system_plugin_state() -> (
+    tuple[list[str], dict[str, Callable[[], type["AbstractBatchSystem"]]]]
+):
     """
     Return a snapshot of the plugin registry that can be restored to remove
     added plugins. Useful for testing the plugin system in-process with other
@@ -190,7 +221,10 @@ def save_batch_system_plugin_state() -> Tuple[List[str], Dict[str, Callable[[],
     snapshot = (list(_registry_keys), dict(_registry))
     return snapshot
 
-def restore_batch_system_plugin_state(snapshot: Tuple[List[str], Dict[str, Callable[[], Type['AbstractBatchSystem']]]]):
+
+def restore_batch_system_plugin_state(
+    snapshot: tuple[list[str], dict[str, Callable[[], type["AbstractBatchSystem"]]]]
+):
     """
     Restore the batch system registry state to a snapshot from
     save_batch_system_plugin_state().