mypy - FSM

src/drunc/fsm/_protocols.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Iterable, List, Optional, Protocol
+
+if TYPE_CHECKING:
+    from druncschema.controller_pb2 import FSMSequence
+
+    from drunc.fsm.core import PreOrPostTransitionSequence
+    from drunc.fsm.transition import Transition
+
+
+class ParameterProtocol(Protocol):
+    name: str
+    value: str
+
+
+class InitConfigurationProtocol(Protocol):
+    parameters: Iterable[ParameterProtocol]
+
+
+class ActorProtocol(Protocol):
+    def get_user_name(self) -> str: ...
+
+
+class DetConfigProtocol(Protocol):
+    id: str
+
+
+class DALProtocol(Protocol):
+    detector_configuration: DetConfigProtocol
+
+
+class DBProtocol(Protocol):
+    def get_dal(self, class_name: str, uid: str) -> DALProtocol: ...
+
+
+class OksKeyProtocol(Protocol):
+    session: str
+
+
+class RuntimeConfigurationProtocol(Protocol):
+    initial_data: str
+    oks_key: OksKeyProtocol
+
+
+class ConfigurationProtocol(Protocol):
+    id: str
+    db: DBProtocol
+    oks_key: OksKeyProtocol
+    initial_data: str
+    parameters: Iterable[ParameterProtocol]
+    name: str
+
+
+class ContextProtocol(Protocol):
+    actor: ActorProtocol
+    configuration: ConfigurationProtocol
+    runinfo: Dict[str, object]
+
+
+class SessionDalProtocol(Protocol):
+    segment: object
+    rte_script: Optional[str]
+
+
+class SSHCommandProtocol(Protocol):
+    def __call__(self, *args: str, _err_to_out: bool = ...) -> object: ...
+
+
+class ShErrorProtocol(Protocol):
+    stdout: bytes
+    stderr: bytes
+
+
+class ActionMethodProtocol(Protocol):
+    __name__: str
+    __module__: str
+    __self__: object
+
+    def __call__(self, *args: object, **kwargs: object) -> object: ...
+
+
+class FSMActionProtocol(Protocol):
+    name: str
+
+
+class FSMxTransitionProtocol(Protocol):
+    transition: str
+    order: list[str]
+    mandatory: list[str]
+
+
+class ConfigProtocol(Protocol):
+    def get_initial_state(self) -> str: ...
+    def get_states(self) -> List[str]: ...
+    def get_transitions(self) -> List[Transition]: ...
+    def get_sequences(self) -> List[FSMSequence]: ...
+    def get_pre_transitions_sequences(
+        self,
+    ) -> Dict[Transition, PreOrPostTransitionSequence]: ...
+    def get_post_transitions_sequences(
+        self,
+    ) -> Dict[Transition, PreOrPostTransitionSequence]: ...
+
+
+class ActionMethod(Protocol):
+    def __call__(self, *args: object, **kwargs: object) -> object: ...

src/drunc/fsm/action_factory.py

@@ -1,9 +1,15 @@
-import inspect
+from __future__ import annotations
 
-import conffwk
+import inspect
+from typing import Dict, Type, cast
 
 import drunc.fsm.exceptions as fsme
 from drunc.exceptions import DruncSetupException
+from drunc.fsm._protocols import (
+    ActionMethodProtocol,
+    ConfigurationProtocol,
+    FSMActionProtocol,
+)
 from drunc.fsm.actions.db_run_registry import DBRunRegistry
 from drunc.fsm.actions.file_logbook import FileLogbook
 from drunc.fsm.actions.file_run_registry import FileRunRegistry
@@ -17,26 +23,34 @@
 
 
 class FSMActionFactory:
-    def __init__(self):
+    _instance: FSMActionFactory | None = None
+
+    def __init__(self) -> None:
         raise DruncSetupException("Call get() instead")
 
-    def _get_pre_transitions(self, action):
-        retr = {}
+    def _get_pre_transitions(
+        self, action: FSMActionProtocol
+    ) -> Dict[str, ActionMethodProtocol]:
+        retr: Dict[str, ActionMethodProtocol] = {}
         for name, method in inspect.getmembers(action):
             if inspect.ismethod(method):
                 if name.startswith("pre_"):
-                    retr[name] = method
+                    retr[name] = cast(ActionMethodProtocol, method)
         return retr
 
-    def _get_post_transitions(self, action):
-        retr = {}
+    def _get_post_transitions(
+        self, action: FSMActionProtocol
+    ) -> Dict[str, ActionMethodProtocol]:
+        retr: Dict[str, ActionMethodProtocol] = {}
         for name, method in inspect.getmembers(action):
             if inspect.ismethod(method):
                 if name.startswith("post_"):
-                    retr[name] = method
+                    retr[name] = cast(ActionMethodProtocol, method)
         return retr
 
-    def _validate_signature(self, name, method, action):
+    def _validate_signature(
+        self, name: str, method: ActionMethodProtocol, action: str
+    ) -> None:
         sig = inspect.signature(method)
 
         if (
@@ -53,7 +67,7 @@ def _validate_signature(self, name, method, action):
             if p.annotation is inspect._empty:
                 raise fsme.MethodSignatureMissingAnnotation(action, name, pname)
 
-    def _validate_action(self, action):
+    def _validate_action(self, action: FSMActionProtocol) -> None:
         pre_transition = self._get_pre_transitions(action)
         post_transition = self._get_post_transitions(action)
 
@@ -67,8 +81,8 @@ def _validate_action(self, action):
             self._validate_signature(k, v, action.name)
 
     def get_action(
-        self, action_name: str, action_configuration: "conffwk.dal.FSMaction"
-    ):
+        self, action_name: str, action_configuration: ConfigurationProtocol
+    ) -> FSMActionProtocol:
         """
         Construct the action interface for the given action name and configuration.
 
@@ -84,7 +98,7 @@ def get_action(
             fsme.InvalidAction: If the constructed action does not have valid pre/post
                 transition methods.
         """
-        iface = None
+        iface: FSMActionProtocol | None = None
         match action_name:
             case "user-provided-run-number":
                 iface = UserProvidedRunNumber(action_configuration)
@@ -116,10 +130,8 @@ def get_action(
 
         return iface
 
-    _instance = None
-
     @classmethod
-    def get(cls):
+    def get(cls: Type[FSMActionFactory]) -> FSMActionFactory:
         if cls._instance is None:
             cls._instance = cls.__new__(cls)
 

src/drunc/fsm/actions/db_run_registry.py

@@ -7,6 +7,7 @@
 from daqconf.jsonify import jsonify_xml_data
 from daqconf.validate import validate_session
 
+from drunc.fsm._protocols import ContextProtocol
 from drunc.fsm.actions.utils import get_dotdrunc_json
 from drunc.fsm.core import FSMAction
 from drunc.fsm.exceptions import (
@@ -20,7 +21,7 @@
 
 
 class DBRunRegistry(FSMAction):
-    def __init__(self, configuration):
+    def __init__(self, configuration: object) -> None:
         super().__init__(name="db-run-registry")
         self.log = get_logger("controller.iface.usvc_db_run_registry")
 
@@ -36,7 +37,12 @@ def __init__(self, configuration):
             ) from exc
         self.timeout = 2
 
-    def pre_start(self, _input_data: dict, _context, **kwargs):
+    def pre_start(
+        self,
+        _input_data: dict[str, object],
+        _context: ContextProtocol,
+        **kwargs: object,
+    ) -> dict[str, object]:
         """
         Upload a copy of the XML and JSON configurations used to take the current run to
         the Run Registry, and insert a new run number entry.
@@ -57,9 +63,7 @@ def pre_start(self, _input_data: dict, _context, **kwargs):
 
         # Seems like run_number isn't in _input_data in post_drain_dataflow so need to
         # initialise it here
-        self.run_number = _input_data[
-            "run"
-        ]
+        self.run_number = _input_data["run"]
 
         # Get the environment variables for upload
         software_version = os.getenv("DUNE_DAQ_BASE_RELEASE")
@@ -119,10 +123,11 @@ def pre_start(self, _input_data: dict, _context, **kwargs):
         with tarfile.open(fileobj=tarball, mode="w:gz") as tar:
             tar.add(xml_filename, arcname=os.path.basename(xml_filename))
             tar.add(json_filename, arcname=os.path.basename(json_filename))
-            tar.add(entry_point_filename, arcname=os.path.basename(entry_point_filename))
+            tar.add(
+                entry_point_filename, arcname=os.path.basename(entry_point_filename)
+            )
         # f_tar.close()
 
-
         # Publish to the run registry
         with open(tarball_name, "rb") as f:
             files = {"file": f}
@@ -162,7 +167,12 @@ def pre_start(self, _input_data: dict, _context, **kwargs):
 
         return _input_data
 
-    def post_drain_dataflow(self, _input_data, _context, **kwargs):
+    def post_drain_dataflow(
+        self,
+        _input_data: dict[str, object],
+        _context: ContextProtocol,
+        **kwargs: object,
+    ) -> None:
         try:
             requests.get(
                 self.API_SOCKET + "/runregistry/updateStopTime/" + str(self.run_number),

src/drunc/fsm/actions/file_logbook.py

@@ -1,18 +1,23 @@
 from typing import Optional
 
+from drunc.fsm._protocols import ConfigurationProtocol, ContextProtocol
 from drunc.fsm.core import FSMAction
 from drunc.utils.utils import now_str
 
 
 class FileLogbook(FSMAction):
-    def __init__(self, configuration):
+    def __init__(self, configuration: ConfigurationProtocol) -> None:
         super().__init__(name="file-logbook")
         self.conf_dict = {p.name: p.value for p in configuration.parameters}
         self.file = self.conf_dict["file_name"]
 
     def post_start(
-        self, _input_data, _context, file_logbook_post: Optional[str] = None, **kwargs
-    ):
+        self,
+        _input_data: dict[str, object],
+        _context: ContextProtocol,
+        file_logbook_post: Optional[str] = None,
+        **kwargs: object,
+    ) -> dict[str, object]:
         with open(self.file, "a") as f:
             f.write(
                 f"Run {_input_data['run']} started by {_context.actor.get_user_name()} at {now_str()}\n"
@@ -24,8 +29,12 @@ def post_start(
         return _input_data
 
     def post_drain_dataflow(
-        self, _input_data, _context, file_logbook_post: str = "", **kwargs
-    ):
+        self,
+        _input_data: dict[str, object],
+        _context: ContextProtocol,
+        file_logbook_post: str = "",
+        **kwargs: object,
+    ) -> dict[str, object]:
         with open(self.file, "a") as f:
             f.write(
                 f"Current run stopped by {_context.actor.get_user_name()} at {now_str()}\n"

src/drunc/fsm/actions/file_run_registry.py

@@ -1,16 +1,23 @@
 import os
+from typing import Dict
 
 from daqconf.consolidate import consolidate_db
 
+from drunc.fsm._protocols import ContextProtocol
 from drunc.fsm.core import FSMAction
 
 
 class FileRunRegistry(FSMAction):
-    def __init__(self, configuration):
+    def __init__(self, configuration: object) -> None:
         super().__init__(name="file-run-registry")
         self.configuration = configuration
 
-    def pre_start(self, _input_data, _context, **kwargs):
+    def pre_start(
+        self,
+        _input_data: Dict[str, object],
+        _context: ContextProtocol,
+        **kwargs: object,
+    ) -> Dict[str, object]:
         run_number = _input_data["run"]
         dest = os.getcwd() + "/run_conf" + str(run_number) + ".data.xml"
         consolidate_db(_context.configuration.initial_data.split(":")[1], f"{dest}")

src/drunc/fsm/actions/some_test_action.py

@@ -1,5 +1,6 @@
 from enum import Enum
 
+from drunc.fsm._protocols import ConfigurationProtocol, ContextProtocol
 from drunc.fsm.core import FSMAction
 
 
@@ -9,18 +10,18 @@ class an_enum(Enum):
 
 
 class SomeTestAction(FSMAction):
-    def __init__(self, configuration):
+    def __init__(self, configuration: ConfigurationProtocol) -> None:
         super().__init__(name="test-action")
 
     def pre_conf(
         self,
-        _input_data: dict,
-        _context,
+        _input_data: dict[str, object],
+        _context: ContextProtocol,
         some_int: int,
         some_str: str,
         some_float: float = 0.2,
-        **kwargs,
-    ) -> dict:
+        **kwargs: object,
+    ) -> dict[str, object]:
         print(f"Running pre_conf of {self.name}")
         _input_data["some_int"] = some_int
         _input_data["some_str"] = some_str