diff --git a/pyproject.toml b/pyproject.toml index be11c1c90..97f26b32f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,57 +1,56 @@ [build-system] -requires = ["setuptools >= 61.0"] build-backend = "setuptools.build_meta" +requires = ["setuptools >= 61.0"] [project] -name = "drunc" description = "A flexible run control infrastructure for a distributed DAQ system" -version = "1.1.4" +name = "drunc" readme = "docs/README.md" requires-python = ">=3.10" +version = "1.1.4" dependencies = [ - "click", - "click-shell", - "protobuf", - "types-protobuf", - "grpcio", - "grpcio-status", - "grpcio-tools", - "types-grpcio", - "gunicorn", - "kafka-python", - "nest-asyncio", - "rich", - "requests", - "Flask", - "Flask-RESTful", - "sh", - "kubernetes", - "pytz", - "psutil", - "paramiko" + "click", + "click-shell", + "protobuf", + "types-protobuf", + "grpcio", + "grpcio-status", + "grpcio-tools", + "types-grpcio", + "gunicorn", + "kafka-python", + "nest-asyncio", + "rich", + "requests", + "Flask", + "Flask-RESTful", + "sh", + "kubernetes", + "pytz", + "psutil", + "paramiko", ] [project.optional-dependencies] +dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75", "types-requests"] prod = ["paramiko[gssapi]"] -dev = ["ruff", "pre-commit", "pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] test = ["pytest", "pytest-cov", "grpcio-testing", "grpcio==1.75", "grpcio-tools==1.75", "grpcio-status==1.75"] [project.scripts] -fake_daq_application = "drunc.apps.fake_daqapp_rest:main" +application-registry-service = "drunc.apps.app_connectivity_server:main" +drunc-check-np0x-cluster = "drunc.apps.check_np0x_cluster:main" +drunc-check-np0x-hw = "drunc.apps.check_np0x_hw_status:main" drunc-controller = "drunc.apps.controller:main" drunc-controller-shell = "drunc.apps.controller_shell:main" +drunc-fsm-tests = "drunc.tests.fsm:main" drunc-process-manager = "drunc.apps.pm:main" drunc-process-manager-shell = "drunc.apps.pm_shell:main" drunc-session-manager = "drunc.apps.session_manager:main" -drunc-unified-shell = "drunc.apps.unified_shell:main" -drunc-fsm-tests = "drunc.tests.fsm:main" -application-registry-service = "drunc.apps.app_connectivity_server:main" -drunc-ssh-doctor = "drunc.apps.ssh_doctor:main" drunc-setup-ssh-config = "drunc.apps.ssh_configurator:main" -drunc-check-np0x-hw = "drunc.apps.check_np0x_hw_status:main" -drunc-check-np0x-cluster = "drunc.apps.check_np0x_cluster:main" - +drunc-ssh-doctor = "drunc.apps.ssh_doctor:main" +drunc-unified-shell = "drunc.apps.unified_shell:main" +fake_daq_application = "drunc.apps.fake_daqapp_rest:main" [tool.setuptools.packages.find] where = ["src"] @@ -61,57 +60,56 @@ where = ["src"] "drunc.data.process_manager.schema" = ["*.json"] [tool.pytest.ini_options] +addopts = "-v --tb=short --cov=drunc --cov=src/drunc" markers = [ - "grpc: marks tests for gRPC isolation (run with --test-grpc)", - "paramiko: marks tests for paramiko isolation (run with --test-paramiko)", + "grpc: marks tests for gRPC isolation (run with --test-grpc)", + "paramiko: marks tests for paramiko isolation (run with --test-paramiko)", ] -addopts = "-v --tb=short --cov=drunc --cov=src/drunc" testpaths = ["tests"] [tool.coverage.run] -source = ["drunc"] omit = ["tests/*"] +source = ["drunc"] # * See https://docs.astral.sh/ruff/rules/ for details on Ruff's linting options [tool.ruff.lint] -select = [ - "E", # pycodestyle errors - "F", # check for errors using PyFlakes - "I", # best practices for import calls - "UP", # suggestions for code modernization - "RUF", # build in Ruff warnings - "R", # refactoring suggestions -] ignore = [ - "E501", # Don't enforce line lengths within a linting context + "E501", # Don't enforce line lengths within a linting context +] +select = [ + "E", # pycodestyle errors + "F", # check for errors using PyFlakes + "I", # best practices for import calls + "UP", # suggestions for code modernization + "RUF", # build in Ruff warnings + "R", # refactoring suggestions ] - [tool.mypy] -disallow_untyped_defs = true -disallow_incomplete_defs = true check_untyped_defs = true +disallow_any_generics = true +disallow_incomplete_defs = true +disallow_subclassing_any = true +disallow_untyped_defs = true no_implicit_optional = true +strict_equality = true +warn_redundant_casts = true warn_return_any = true warn_unused_ignores = true -warn_redundant_casts = true -strict_equality = true -disallow_any_generics = true -disallow_subclassing_any = true +disallow_any_decorated = true disallow_any_explicit = true disallow_any_expr = false -disallow_any_decorated = true disallow_any_unimported = true -warn_unused_configs = true show_error_codes = true +warn_unused_configs = true # These overrides are because the library stubs dont exist and not on typeshed [[tool.mypy.overrides]] -module = ["google.rpc.*"] ignore_missing_imports = true +module = ["google.rpc.*"] [[tool.mypy.overrides]] +ignore_missing_imports = true module = ["conffwk"] -ignore_missing_imports = true \ No newline at end of file diff --git a/src/drunc/fsm/_protocols.py b/src/drunc/fsm/_protocols.py new file mode 100644 index 000000000..d2431b58a --- /dev/null +++ b/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: ... diff --git a/src/drunc/fsm/action_factory.py b/src/drunc/fsm/action_factory.py index 5d23a1218..d4024f075 100644 --- a/src/drunc/fsm/action_factory.py +++ b/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) diff --git a/src/drunc/fsm/actions/db_run_registry.py b/src/drunc/fsm/actions/db_run_registry.py index 79b8e6152..3767b07b7 100644 --- a/src/drunc/fsm/actions/db_run_registry.py +++ b/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), diff --git a/src/drunc/fsm/actions/file_logbook.py b/src/drunc/fsm/actions/file_logbook.py index b5cd530a9..f776160e6 100644 --- a/src/drunc/fsm/actions/file_logbook.py +++ b/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" diff --git a/src/drunc/fsm/actions/file_run_registry.py b/src/drunc/fsm/actions/file_run_registry.py index 4254f0a8b..b5b7cd145 100644 --- a/src/drunc/fsm/actions/file_run_registry.py +++ b/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}") diff --git a/src/drunc/fsm/actions/some_test_action.py b/src/drunc/fsm/actions/some_test_action.py index 669b5c44f..2a42e1de0 100644 --- a/src/drunc/fsm/actions/some_test_action.py +++ b/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 diff --git a/src/drunc/fsm/actions/thread_pinning.py b/src/drunc/fsm/actions/thread_pinning.py index b1fa976e5..0dd5323d1 100644 --- a/src/drunc/fsm/actions/thread_pinning.py +++ b/src/drunc/fsm/actions/thread_pinning.py @@ -2,9 +2,10 @@ from os import environ import conffwk -from sh import Command, ErrorReturnCode +from sh import Command, ErrorReturnCode # type: ignore[import-untyped] from drunc.exceptions import DruncSetupException +from drunc.fsm._protocols import ContextProtocol, InitConfigurationProtocol from drunc.fsm.core import FSMAction from drunc.fsm.exceptions import ThreadPinningFailed from drunc.process_manager.oks_parser import collect_apps @@ -13,12 +14,14 @@ class ThreadPinning(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: InitConfigurationProtocol) -> None: super().__init__(name="thread-pinning") self.log = get_logger("controller.iface.thread-pinning") self.conf_dict = {p.name: p.value for p in configuration.parameters} - def pin_thread(self, thread_pinning_file, configuration, session): + def pin_thread( + self, thread_pinning_file: str, configuration: str, session: str + ) -> None: db = conffwk.Configuration(configuration) session_dal = db.get_dal(class_name="Session", uid=session) @@ -27,7 +30,7 @@ def pin_thread(self, thread_pinning_file, configuration, session): session_name=session, session_dal_obj=session_dal, segment_obj=session_dal.segment, - env=environ, + env=dict(environ), tree_prefix=[], ) @@ -86,7 +89,12 @@ def pin_thread(self, thread_pinning_file, configuration, session): if failed_hosts: raise ThreadPinningFailed(failed_hosts_error_str) - def post_conf(self, _input_data, _context, **kwargs): + def post_conf( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, + ) -> dict[str, object]: if "post_conf" in self.conf_dict: self.pin_thread( self.conf_dict["post_conf"], @@ -95,7 +103,12 @@ def post_conf(self, _input_data, _context, **kwargs): ) return _input_data - def post_start(self, _input_data, _context, **kwargs): + def post_start( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, + ) -> dict[str, object]: if "post_start" in self.conf_dict: self.pin_thread( self.conf_dict["post_start"], @@ -104,7 +117,12 @@ def post_start(self, _input_data, _context, **kwargs): ) return _input_data - def pre_conf(self, _input_data, _context, **kwargs): + def pre_conf( + self, + _input_data: dict[str, object], + _context: ContextProtocol, + **kwargs: object, + ) -> dict[str, object]: if "pre_conf" in self.conf_dict: self.pin_thread( self.conf_dict["pre_conf"], diff --git a/src/drunc/fsm/actions/timing/master_send_fl_command.py b/src/drunc/fsm/actions/timing/master_send_fl_command.py index a289af2a2..ee7858316 100644 --- a/src/drunc/fsm/actions/timing/master_send_fl_command.py +++ b/src/drunc/fsm/actions/timing/master_send_fl_command.py @@ -1,19 +1,20 @@ +from drunc.fsm._protocols import ContextProtocol from drunc.fsm.core import FSMAction class MasterSendFLCommand(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: super().__init__(name="master-send-fl-command") def pre_master_send_fl_command( self, - _input_data, - _context, + _input_data: dict[str, object], + _context: ContextProtocol, fl_cmd_id: int, channel: int, number_of_commands_to_send: int, - **kwargs, - ): + **kwargs: object, + ) -> dict[str, object]: # parse fl_cmd_id... _input_data["fl_cmd_id"] = fl_cmd_id _input_data["channel"] = channel diff --git a/src/drunc/fsm/actions/trigger_rate_specifier.py b/src/drunc/fsm/actions/trigger_rate_specifier.py index 189da9002..cd05cf10a 100644 --- a/src/drunc/fsm/actions/trigger_rate_specifier.py +++ b/src/drunc/fsm/actions/trigger_rate_specifier.py @@ -1,12 +1,17 @@ +from drunc.fsm._protocols import ConfigurationProtocol, ContextProtocol from drunc.fsm.core import FSMAction class TriggerRateSpecifier(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: ConfigurationProtocol) -> None: super().__init__(name="trigger-rate-specifier") def pre_change_rate( - self, _input_data: dict, _context, trigger_rate: float, **kwargs - ): + self, + _input_data: dict[str, object], + _context: ContextProtocol, + trigger_rate: float, + **kwargs: object, + ) -> dict[str, object]: _input_data["trigger_rate"] = trigger_rate return _input_data diff --git a/src/drunc/fsm/actions/user_provided_run_number.py b/src/drunc/fsm/actions/user_provided_run_number.py index ea97dec5f..dba128c22 100644 --- a/src/drunc/fsm/actions/user_provided_run_number.py +++ b/src/drunc/fsm/actions/user_provided_run_number.py @@ -6,20 +6,22 @@ class UserProvidedRunNumber(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: super().__init__(name="run-number") def pre_start( self, - _input_data: dict, - _context, + _input_data: dict[str, object], + _context: object, run_number: int, run_type: Optional[str] = "TEST", disable_data_storage: bool = False, trigger_rate: Optional[float] = None, - **kwargs, - ): - run_type = validate_run_type(run_type.upper()) + **kwargs: object, + ) -> dict[str, object]: + + safe_run_type = run_type.upper() if run_type is not None else "TEST" + run_type = validate_run_type(safe_run_type) _input_data["production_vs_test"] = run_type _input_data["run"] = run_number _input_data["disable_data_storage"] = disable_data_storage diff --git a/src/drunc/fsm/actions/usvc_elisa_logbook.py b/src/drunc/fsm/actions/usvc_elisa_logbook.py index 7dfcb4716..2b2d747d5 100644 --- a/src/drunc/fsm/actions/usvc_elisa_logbook.py +++ b/src/drunc/fsm/actions/usvc_elisa_logbook.py @@ -4,6 +4,7 @@ import requests +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 ( @@ -15,7 +16,7 @@ class ElisaLogbook(FSMAction): - def __init__(self): + def __init__(self) -> None: super().__init__(name="elisa-logbook") self.log = get_logger("controller.iface.elisa-logbook") @@ -65,7 +66,6 @@ def __init__(self): ) self.log.warning(warn_msg) else: - warn_msg: str = "" if default_elisa_logbook: warn_msg = ( "You need to update your ~/.drunc.json: The default ELisA logbook " @@ -83,10 +83,14 @@ def __init__(self): self.timeout = 5 def post_start( - self, _input_data: dict, _context, elisa_post: Optional[str] = None, **kwargs - ): + self, + _input_data: dict[str, object], + _context: ContextProtocol, + elisa_post: Optional[str] = None, + **kwargs: object, + ) -> dict[str, object]: if self.elisa_hardware in self.no_publish_hardware: - return + return {} text = "" self.thread_id = None # Clear this value here, so that if it fails stop can't reply to an old message @@ -142,10 +146,14 @@ def post_start( return _input_data def post_drain_dataflow( - self, _input_data, _context, elisa_post: Optional[str] = None, **kwargs - ): + self, + _input_data: dict[str, object], + _context: ContextProtocol, + elisa_post: Optional[str] = None, + **kwargs: object, + ) -> dict[str, object]: if self.elisa_hardware in self.no_publish_hardware: - return + return {} text = "" if elisa_post is not None: self.log.info( diff --git a/src/drunc/fsm/actions/usvc_provided_run_number.py b/src/drunc/fsm/actions/usvc_provided_run_number.py index a3553c966..d63ef072b 100644 --- a/src/drunc/fsm/actions/usvc_provided_run_number.py +++ b/src/drunc/fsm/actions/usvc_provided_run_number.py @@ -1,8 +1,9 @@ import time -from typing import Optional +from typing import Optional, cast import requests +from drunc.fsm._protocols import ContextProtocol from drunc.fsm.actions.utils import get_dotdrunc_json, validate_run_type from drunc.fsm.core import FSMAction from drunc.fsm.exceptions import CannotGetRunNumber, DotDruncJsonIncorrectFormat @@ -10,7 +11,7 @@ class UsvcProvidedRunNumber(FSMAction): - def __init__(self, configuration): + def __init__(self, configuration: object) -> None: self.log = get_logger("controller.iface.usvc_run_number") super().__init__(name="usvc-provided-run-number") dotdrunc = get_dotdrunc_json() @@ -27,13 +28,13 @@ def __init__(self, configuration): def pre_start( self, - _input_data: dict, - _context, + _input_data: dict[str, object], + _context: ContextProtocol, run_type: str, disable_data_storage: bool = False, trigger_rate: Optional[float] = None, - **kwargs, - ): + **kwargs: object, + ) -> dict[str, object]: run_type = validate_run_type(run_type.upper()) _input_data["production_vs_test"] = run_type _input_data["run"] = self._getnew_run_number() @@ -45,7 +46,7 @@ def pre_start( return _input_data - def _getnew_run_number(self): + def _getnew_run_number(self) -> int: try: req = requests.get( self.API_SOCKET + "/runnumber/getnew", @@ -68,5 +69,6 @@ def _getnew_run_number(self): self.log.error(error) raise CannotGetRunNumber(error) from exc - self.run = req.json()[0][0][0] + self.run = cast(int, req.json()[0][0][0]) + return self.run diff --git a/src/drunc/fsm/actions/utils.py b/src/drunc/fsm/actions/utils.py index b32d5eed3..e9e3efe52 100644 --- a/src/drunc/fsm/actions/utils.py +++ b/src/drunc/fsm/actions/utils.py @@ -11,8 +11,10 @@ def validate_run_type(run_type: str) -> str: """Validate the run type - :param run_type: the run type - :return: the validated run type + Args: + run_type: the run type + Returns: + the validated run type """ RUN_TYPES = ["PROD", "TEST"] if run_type not in RUN_TYPES: @@ -22,7 +24,7 @@ def validate_run_type(run_type: str) -> str: return run_type -def get_dotdrunc_json(path: str | None = None): +def get_dotdrunc_json(path: str | None = None): # type: ignore[no-untyped-def] # Resolution order: DOTDRUNC env var -> provided path -> default path file_path = os.getenv("DOTDRUNC") or path or "~/.drunc.json" try: diff --git a/src/drunc/fsm/configuration.py b/src/drunc/fsm/configuration.py index 16b3890bf..6ef108be5 100644 --- a/src/drunc/fsm/configuration.py +++ b/src/drunc/fsm/configuration.py @@ -1,6 +1,11 @@ -import conffwk +from __future__ import annotations + +from typing import List, cast + +from conffwk.dal import FSMData, FSMxTransition from druncschema.controller_pb2 import FSMSequence +from drunc.fsm._protocols import FSMActionProtocol from drunc.fsm.action_factory import FSMActionFactory from drunc.fsm.core import PreOrPostTransitionSequence from drunc.fsm.transition import Transition @@ -11,11 +16,13 @@ class FSMConfHandler(ConfHandler): """Handler for FSM configuration.""" + data: FSMData + def _fill_pre_post_transition_sequence_oks( self, prefix: str, transition: Transition, - data: list["conffwk.dal.FSMxTransition"], + data: List[FSMxTransition] | None, ) -> PreOrPostTransitionSequence: """ Fill the pre or post transition sequence for a given transition. @@ -72,14 +79,14 @@ def _post_process_oks(self) -> None: Raises: None """ - raw = self._raw_data + raw = cast(FSMData, self._raw_data) # Define the data structures to store the FSM configuration self.log.debug("_post_process_oks configuration") - self._pre_transitions: dict[Transition, PreOrPostTransitionSequence] = {} - self._post_transitions: dict[Transition, PreOrPostTransitionSequence] = {} + self.pre_transitions: dict[Transition, PreOrPostTransitionSequence] = {} + self.post_transitions: dict[Transition, PreOrPostTransitionSequence] = {} self.actions: dict[ - str, "conffwk.dal.FSMAction" + str, FSMActionProtocol ] = {} # (e.g. "thread_pinning": thread_pinning FSM action object) self.transitions: list[Transition] = [] self.sequences: list[FSMSequence] = [] @@ -88,12 +95,12 @@ def _post_process_oks(self) -> None: # Fill the actions dictionary with the FSMAction objects corresponding to the # action names defined in the configuration - for action in raw.actions: # type: 'conffwk.dal.FSMAction' + for action in raw.actions: self.actions[action.id] = FSMActionFactory.get().get_action( action.id, action ) - for transition in raw.transitions: # type: 'conffwk.dal.FSMTransition' + for transition in raw.transitions: tr = Transition( name=transition.id, source=transition.source, @@ -114,12 +121,13 @@ def _post_process_oks(self) -> None: ) # Add the pre and post transition sequence arguments to the transition - tr.arguments += pre_transitions.get_arguments() - tr.arguments += post_transitions.get_arguments() + if tr.arguments is not None: + tr.arguments += pre_transitions.get_arguments() + tr.arguments += post_transitions.get_arguments() # Store the pre and post transition sequences for the transition - self._pre_transitions[tr] = pre_transitions - self._post_transitions[tr] = post_transitions + self.pre_transitions[tr] = pre_transitions + self.post_transitions[tr] = post_transitions # Add the transition to the list of transitions self.transitions += [tr] @@ -130,10 +138,10 @@ def _post_process_oks(self) -> None: cmd_ids = [cmd.id for cmd in sequence.sequence] self.sequences.append(FSMSequence(id=seq_id, command_ids=cmd_ids)) - def get_actions(self) -> dict[str, "conffwk.dal.FSMAction"]: + def get_actions(self) -> dict[str, FSMActionProtocol]: return self.actions - def get_initial_state(self) -> str | None: + def get_initial_state(self) -> str: return self.initial_state def get_states(self) -> list[str]: @@ -145,12 +153,12 @@ def get_transitions(self) -> list[Transition]: def get_pre_transitions_sequences( self, ) -> dict[Transition, PreOrPostTransitionSequence]: - return self._pre_transitions + return self.pre_transitions def get_post_transitions_sequences( self, ) -> dict[Transition, PreOrPostTransitionSequence]: - return self._post_transitions + return self.post_transitions def get_sequences(self) -> list[FSMSequence]: return self.sequences diff --git a/src/drunc/fsm/core.py b/src/drunc/fsm/core.py index 7ebb2c7e5..8d5b0e2e0 100644 --- a/src/drunc/fsm/core.py +++ b/src/drunc/fsm/core.py @@ -1,14 +1,27 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Dict, Optional, Union, cast + +if TYPE_CHECKING: + from drunc.fsm._protocols import ( + ActionMethodProtocol, + ConfigProtocol, + ContextProtocol, + FSMActionProtocol, + ) + + # Define the abcs first to avoid circular imports class FSMAction: """Abstract class defining a generic action""" - def __init__(self, name): + def __init__(self, name: str) -> None: self.name = name class Callback: - def __init__(self, method: "conffwk.dal.FSMAction", mandatory: bool = True): - self.method: "conffwk.dal.FSMAction" = method + def __init__(self, method: ActionMethodProtocol, mandatory: bool = True) -> None: + self.method: ActionMethodProtocol = method self.mandatory: bool = mandatory @@ -17,9 +30,7 @@ def __init__(self, method: "conffwk.dal.FSMAction", mandatory: bool = True): from dataclasses import dataclass from enum import Enum from inspect import Parameter, signature -from typing import Optional, Union -import conffwk from druncschema.controller_pb2 import Argument, FSMSequence from druncschema.generic_pb2 import bool_msg, float_msg, int_msg, string_msg from google.protobuf import any_pb2 @@ -32,7 +43,7 @@ def __init__(self, method: "conffwk.dal.FSMAction", mandatory: bool = True): class PreOrPostTransitionSequence: - def __init__(self, transition: Transition, pre_or_post: str = "pre"): + def __init__(self, transition: Transition, pre_or_post: str = "pre") -> None: self.transition: Transition = transition if pre_or_post not in ["pre", "post"]: raise DruncSetupException( @@ -41,12 +52,10 @@ def __init__(self, transition: Transition, pre_or_post: str = "pre"): self.prefix: str = pre_or_post - self.sequence: list(Callback) = [] + self.sequence: list[Callback] = [] self.log = get_logger("controller.core.PreOrPostTransitionSequence") - def add_callback( - self, action: "conffwk.dal.FSMaction", mandatory: bool = True - ) -> None: + def add_callback(self, action: FSMActionProtocol, mandatory: bool = True) -> None: """ Add a callback to the sequence. The method to be called will be determined by the name of the transition and the prefix (pre or post). @@ -83,7 +92,7 @@ def add_callback( ) ] - def __str__(self): + def __str__(self) -> str: return ", ".join( [ f"{cb.method.__self__.__class__.__name__} (mandatory={cb.mandatory})" @@ -91,26 +100,34 @@ def __str__(self): ] ) - def execute(self, transition_data, transition_args, ctx=None): + def execute( + self, + transition_data: str | None, + transition_args: Dict[str, object], + ctx: ContextProtocol, + ) -> str: self.log.debug(f"{transition_data=}, {transition_args=}") if not transition_data: transition_data = "{}" try: - input_data = json.loads(transition_data) + input_data = cast(dict[str, object], json.loads(transition_data)) except: raise fsme.TransitionDataOfIncorrectFormat(transition_data) for callback in self.sequence: try: - self.log.debug(f"data before callback: {input_data}") - self.log.debug( + self.log.critical(f"data before callback: {input_data}") + self.log.critical( f"executing the callback: {callback.method.__name__} from {callback.method.__module__}" ) - input_data = callback.method( - _input_data=input_data, _context=ctx, **transition_args + input_data = cast( + dict[str, object], + callback.method( + _input_data=input_data, _context=ctx, **transition_args + ), ) - self.log.debug(f"data after callback: {input_data}") + self.log.critical(f"data after callback: {input_data}") if input_data: ctx.runinfo.update(input_data) @@ -149,11 +166,11 @@ def get_arguments(self) -> list[Argument]: """ # Construct the list of arguments by looking at the signature of the methods - arguments: list(Argument) = [] + arguments: list[Argument] = [] # Check that there are no duplicate parameter names across the callbacks # otherwise, we won't know which one to use when executing the sequence - all_sequence_arguments: set(str) = set() # set(Argument names) + all_sequence_arguments: set[str] = set() # set(Argument names) # Iterate over the callbacks, construct the list of arguments for callback in self.sequence: @@ -184,7 +201,7 @@ def get_arguments(self) -> list[Argument]: # Determine the type of the argument, and set the default value if it is # optional. If the type is not one of the supported types, raise an # error - t: Argument.Type = Argument.Type.INT + t: int = Argument.Type.INT if p.annotation in (str, Optional[str], Union[str, None]): t = Argument.Type.STRING @@ -256,7 +273,7 @@ class FSMDestinationResult: class FSM: - def __init__(self, conf): + def __init__(self, conf: ConfigProtocol) -> None: self.log = get_logger("controller.core.FSM") self.configuration = conf @@ -280,7 +297,7 @@ def __init__(self, conf): self.log.debug(f"Pre transition: {self.pre_transition_sequences[t]}") self.log.debug(f"Post transition: {self.post_transition_sequences[t]}") - def _enusure_unique_transition(self, transitions): + def _enusure_unique_transition(self, transitions: list[Transition]) -> None: a_set = set() for t in transitions: if t.name in a_set: @@ -299,10 +316,14 @@ def get_all_sequences(self) -> list[FSMSequence]: """Grab all the transitions""" return self.sequences - def is_destination_of_this_transition(self, state, transition) -> bool: - return transition.destination == state + def is_destination_of_this_transition( + self, state: str, transition: Transition + ) -> bool: + return bool(transition.destination == state) - def get_destination_state(self, source_state, transition) -> FSMDestinationResult: + def get_destination_state( + self, source_state: str, transition: Transition + ) -> FSMDestinationResult: """Tells us where a particular transition will take us, given the source_state""" right_name = [t for t in self.transitions if t == transition] @@ -336,7 +357,7 @@ def get_destination_state(self, source_state, transition) -> FSMDestinationResul destination_type=FSMDestinationType.TRANSITION_NOT_VALID, ) - def get_executable_transitions(self, source_state) -> list[Transition]: + def get_executable_transitions(self, source_state: str) -> list[Transition]: valid_transitions = [] for tr in self.transitions: @@ -349,7 +370,7 @@ def get_executable_transitions(self, source_state) -> list[Transition]: return valid_transitions - def get_executable_sequences(self, source_state) -> list[FSMSequence]: + def get_executable_sequences(self, source_state: str) -> list[FSMSequence]: valid_sequences = [] for seq in self.sequences: @@ -366,7 +387,7 @@ def get_executable_sequences(self, source_state) -> list[FSMSequence]: return valid_sequences - def get_transition(self, transition_name) -> Transition: + def get_transition(self, transition_name: str) -> Transition: self.log.debug(f"Searching for transition {transition_name}") transition = [t for t in self.transitions if t.name == transition_name] self.log.debug(f"Found transition {transition}") @@ -374,22 +395,30 @@ def get_transition(self, transition_name) -> Transition: raise fsme.NoTransitionOfName(transition_name) return transition[0] - def can_execute_transition(self, source_state, transition) -> bool: + def can_execute_transition(self, source_state: str, transition: Transition) -> bool: """Check that this transition is allowed given the source_state""" self.log.debug(f"can_execute_transition {transition.source!s} {source_state}") - return regex_match(transition.source, source_state) + return bool(regex_match(transition.source, source_state)) def prepare_transition( - self, transition, transition_data, transition_args, ctx=None - ): + self, + transition: Transition, + transition_data: str, + transition_args: Dict[str, object], + ctx: ContextProtocol, + ) -> str: transition_data = self.pre_transition_sequences[transition].execute( transition_data, transition_args, ctx ) return transition_data def finalise_transition( - self, transition, transition_data, transition_args, ctx=None - ): + self, + transition: Transition, + transition_data: str, + transition_args: Dict[str, object], + ctx: ContextProtocol, + ) -> str: transition_data = self.post_transition_sequences[transition].execute( transition_data, transition_args, ctx ) diff --git a/src/drunc/fsm/exceptions.py b/src/drunc/fsm/exceptions.py index 66148eaa5..acffab9ee 100644 --- a/src/drunc/fsm/exceptions.py +++ b/src/drunc/fsm/exceptions.py @@ -6,15 +6,15 @@ class FSMException(DruncCommandException): class NoTransitionOfName(FSMException): - def __init__(self, transition_name): + def __init__(self, transition_name: str) -> None: self.message = f'Transition "{transition_name}" does not exist' - super(NoTransitionOfName, self).__init__(self.message) + super().__init__(self.message) class DuplicateTransition(FSMException): """If a transition has the same name as another""" - def __init__(self, transition_name): + def __init__(self, transition_name: str) -> None: self.message = f'Transition "{transition_name}" is a duplicate' super().__init__(self.message) @@ -22,7 +22,7 @@ def __init__(self, transition_name): class InvalidTransition(FSMException): """Raised when the transition isn't in the list of currently accessible transitions""" - def __init__(self, transition, state): + def __init__(self, transition: str, state: str) -> None: self.message = f"Transition {transition} is not allowed from the state {state}." super().__init__(self.message) @@ -30,7 +30,7 @@ def __init__(self, transition, state): class UnregisteredTransition(FSMException): """Raised when the transition is allowed, but we can't find the implementation""" - def __init__(self, transition): + def __init__(self, transition: str) -> None: self.message = f"Implementation of {transition} not found." super().__init__(self.message) @@ -38,7 +38,7 @@ def __init__(self, transition): class UnknownAction(FSMException): """Raised when a plugin name is provided that does not correspond to any files in /plugins""" - def __init__(self, name): + def __init__(self, name: str) -> None: self.message = f'"{name}" is not a known plugin.' super().__init__(self.message) @@ -46,7 +46,7 @@ def __init__(self, name): class MissingArgument(FSMException): """Raised when a mandatory argument is not provided for a transition""" - def __init__(self, param, name): + def __init__(self, param: str, name: str) -> None: self.message = f'The mandatory argument "{param}" was not provided to the transition {name}' super().__init__(self.message) @@ -54,7 +54,7 @@ def __init__(self, param, name): class MissingArgumentValue(FSMException): """Raised when a mandatory argument is not provided for a transition""" - def __init__(self): + def __init__(self) -> None: self.message = "A passed argument does not have an associated value, arguments are key-value pairs." super().__init__(self.message) @@ -62,13 +62,13 @@ def __init__(self): class DoubleArgument(FSMException): """Raised when an argument is provided more than once""" - def __init__(self, txt): + def __init__(self, txt: str) -> None: self.message = txt super().__init__(self.message) class UnhandledArgumentType(FSMException): - def __init__(self, annotation): + def __init__(self, annotation: str) -> None: self.message = f'The argument "{annotation}" cannot be handled' super().__init__(self.message) @@ -76,7 +76,7 @@ def __init__(self, annotation): class UnknownArgument(FSMException): """Raised when an unwanted argument is given to a transition""" - def __init__(self, param, name): + def __init__(self, param: str, name: str) -> None: self.message = ( f'The mandatory argument "{param}" is not required by transition {name}' ) @@ -86,7 +86,7 @@ def __init__(self, param, name): class InvalidAction(FSMException): """Raised when an action doesn't have pre/post transitions""" - def __init__(self, iface): + def __init__(self, iface: str) -> None: self.message = ( f'The action "{iface}" does not have any pre or post transition method' ) @@ -96,48 +96,49 @@ def __init__(self, iface): class InvalidActionMethod(FSMException): """Raised when an action doesn't have the pre/post transitions arguments""" - def __init__(self, iface, method): + def __init__(self, iface: str, method: str) -> None: self.message = f'The action "{iface}" method {method} does not have the correct arguements, each one should have at least "_input_data" and "**kwargs", and have type annotations' super().__init__(self.message) class MethodSignatureMissingAnnotation(FSMException): - def __init__(self, iface, method, pname): + def __init__(self, iface: str, method: str, pname: str) -> None: self.message = f'The action "{iface}" method {method} does not have the correct arguement annotation for "{pname}", provide "argument:int" to your pre/post methods.' super().__init__(self.message) class TransitionDataOfIncorrectFormat(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f'The data "{data}" could not be interpreted as json' - super(MethodSignatureMissingAnnotation, self).__init__(self.message) + super().__init__(self.message) class CannotGetRunNumber(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Could not get Run Number because {data}" super().__init__(self.message) + class DBRunRegistryConfigurationError(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"RunRegistryDB configuration error: {data}" super().__init__(self.message) class CannotInsertRunNumber(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Could not insert Run into RunRegistryDB because {data}" super().__init__(self.message) class CannotUpdateStopTime(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Could not update stop time in RunRegistryDB because {data}" super().__init__(self.message) class InvalidDataReturnByFSMAction(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = ( f"The action returns an incorrect object which isn't serialisable: {data}" ) @@ -145,19 +146,19 @@ def __init__(self, data): class ThreadPinningFailed(FSMException): - def __init__(self, host): + def __init__(self, host: str) -> None: self.message = f'The thread pinning on "{host}" failed' super().__init__(self.message) class CannotGetSoftwareVersion(FSMException): - def __init__(self): + def __init__(self) -> None: self.message = "RunRegistryDB: dunedaq version not in the variable env DUNE_DAQ_BASE_RELEASE! Exit drunc and export DUNE_DAQ_BASE_RELEASE=dunedaq-vX.XX.XX\n" super().__init__(self.message) class CannotSendElisaMessage(FSMException): - def __init__(self, data): + def __init__(self, data: object) -> None: self.message = f"Cannot send message to ELisA because {data}. Do it manually at https://np-vd-coldbox-elog.app.cern.ch or https://pdsp-elog.app.cern.ch!" super().__init__(self.message) diff --git a/src/drunc/fsm/transition.py b/src/drunc/fsm/transition.py index 4faafba37..cae2a6ab1 100644 --- a/src/drunc/fsm/transition.py +++ b/src/drunc/fsm/transition.py @@ -1,12 +1,26 @@ +from __future__ import annotations + +from typing import List, Optional + +from druncschema.controller_pb2 import Argument + + class Transition: - def __init__(self, name, source, destination, arguments=[], help: str = ""): + def __init__( + self, + name: str, + source: str, + destination: str, + arguments: Optional[List[Argument]] = None, + help: str = "", + ) -> None: self.source = source self.destination = destination self.name = name self.arguments = arguments self.help = help - def __eq__(self, another): + def __eq__(self, another: object) -> bool: same_name = hasattr(another, "name") and self.name == another.name same_destination = ( hasattr(another, "destination") and self.destination == another.destination @@ -14,8 +28,8 @@ def __eq__(self, another): same_source = hasattr(another, "source") and self.source == another.source return same_name and same_destination and same_source - def __hash__(self): + def __hash__(self) -> int: return hash(self.__str__()) - def __str__(self): + def __str__(self) -> str: return f'"{self.name}": "{self.source}" → "{self.destination}"' diff --git a/src/drunc/fsm/utils.py b/src/drunc/fsm/utils.py index 6c8c0c44c..80f7e26d6 100644 --- a/src/drunc/fsm/utils.py +++ b/src/drunc/fsm/utils.py @@ -1,4 +1,4 @@ -from typing import MutableMapping +from typing import MutableMapping, Optional from druncschema.controller_pb2 import ( Argument, @@ -6,7 +6,9 @@ FSMCommandsDescription, ) from druncschema.generic_pb2 import bool_msg, float_msg, int_msg, string_msg -from google.protobuf.any_pb2 import Any + +# prevent confusion with typing.Any +from google.protobuf.any_pb2 import Any as PbAny import drunc.fsm.exceptions as fsme from drunc.fsm.transition import Transition @@ -21,14 +23,14 @@ def convert_fsm_transition(transitions: list[Transition]) -> FSMCommandsDescript transitions (list[Transition]): A list of FSM transitions. Returns: - FSMCommandsDescription: A FSMCommandsDescription containing the converted + FSMCommandsDescription: A FSMCommandsDescription containing the converted transitions. - + Raises: None. """ commands_description = FSMCommandsDescription() - for t in transitions: # type: Transition + for t in transitions: commands_description.commands.append( FSMCommandDescription( name=t.name, @@ -41,16 +43,18 @@ def convert_fsm_transition(transitions: list[Transition]) -> FSMCommandsDescript return commands_description -def decode_fsm_arguments(arguments: MutableMapping[str, Any], arguments_format: list[Argument]) -> dict[str, str | int | float | bool]: +def decode_fsm_arguments( + arguments: MutableMapping[str, PbAny], arguments_format: list[Argument] +) -> dict[str, str | int | float | bool]: """ Decodes the arguments of a FSM command. - Note there is separate logic to validate whether the required arguments are all + Note there is separate logic to validate whether the required arguments are all present at the click.core.Command level, this is a safeguard to support the multiple drunc operating modes. Args: - arguments (MutableMapping[str, Any]): The arguments to decode. + arguments (MutableMapping[str, PbAny]): The arguments to decode. arguments_format (list[Argument]): The format of the arguments. Returns: @@ -61,15 +65,18 @@ def decode_fsm_arguments(arguments: MutableMapping[str, Any], arguments_format: fsme.UnhandledArgumentType: If an argument type is not handled. """ - def get_argument(name, arguments): - for n, k in arguments.items(): + # Added explicit type hints to the nested function + def get_argument(name: str, args: MutableMapping[str, PbAny]) -> Optional[PbAny]: + for n, k in args.items(): if n == name: return k return None - out_dict = {} + # Explicitly type out_dict + out_dict: dict[str, str | int | float | bool] = {} + for arg in arguments_format: - arg_value = get_argument(arg.name, arguments) + arg_value: Optional[PbAny] = get_argument(arg.name, arguments) if arg.presence == Argument.Presence.MANDATORY and arg_value is None: raise fsme.MissingArgument(arg.name, "") @@ -87,5 +94,6 @@ def get_argument(name, arguments): case Argument.Type.BOOL: out_dict[arg.name] = unpack_any(arg_value, bool_msg).value case _: - raise fsme.UnhandledArgumentType(arg.type) + raise fsme.UnhandledArgumentType(str(arg.type)) + return out_dict diff --git a/src/drunc/utils/grpc_utils.py b/src/drunc/utils/grpc_utils.py index a228de8a7..6741ebb67 100644 --- a/src/drunc/utils/grpc_utils.py +++ b/src/drunc/utils/grpc_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Callable, NoReturn, cast +from typing import Callable, NoReturn, TypeVar, cast import grpc from druncschema.generic_pb2 import PlainText @@ -76,7 +76,10 @@ def pack_to_any(data: Message) -> any_pb2.Any: return any -def unpack_any(data: any_pb2.Any, format: type[Message]) -> Message: +T = TypeVar("T") + + +def unpack_any(data: any_pb2.Any, format: type[T]) -> T: """Unpack an Any message into a specific protobuf format. Args: