misc: pre-commit / uv.lock

Explicitly declare protocol implementations

src/plumpy/broadcast_filter.py

@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+# type: ignore
 import re
 import typing
 

src/plumpy/controller.py

@@ -2,7 +2,7 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Any, Hashable, Optional, Protocol, Union
+from typing import Any, Hashable, Optional, Protocol, Union, runtime_checkable
 
 from plumpy import loaders
 from plumpy.message import MessageType
@@ -12,6 +12,7 @@
 ProcessStatus = Any
 
 
+@runtime_checkable
 class ProcessController(Protocol):
     """
     Control processes using coroutines that will send messages and wait
@@ -26,7 +27,7 @@ def get_status(self, pid: 'PID_TYPE') -> ProcessStatus:
         """
         ...
 
-    def pause_process(self, pid: 'PID_TYPE', msg: str | None = None) -> ProcessResult:
+    def pause_process(self, pid: 'PID_TYPE', msg_text: str | None = None) -> Any:
         """
         Pause the process
 
@@ -52,8 +53,7 @@ def play_process(self, pid: 'PID_TYPE') -> ProcessResult:
         ...
 
     def play_all(self) -> None:
-        """Play all processes that are subscribed to the same coordinator
-        """
+        """Play all processes that are subscribed to the same coordinator"""
 
     def kill_process(self, pid: 'PID_TYPE', msg_text: str | None = None) -> Any:
         """Kill the process

src/plumpy/coordinator.py

@@ -1,14 +1,15 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Callable, Hashable, Protocol
 from re import Pattern
+from typing import TYPE_CHECKING, Any, Callable, Hashable, Protocol, runtime_checkable
 
 if TYPE_CHECKING:
     ID_TYPE = Hashable
     Receiver = Callable[..., Any]
 
 
+@runtime_checkable
 class Coordinator(Protocol):
     def hook_rpc_receiver(
         self,
@@ -36,7 +37,11 @@ def unhook_broadcast_receiver(self, identifier: 'ID_TYPE | None') -> None: ...
 
     def unhook_task_receiver(self, identifier: 'ID_TYPE') -> None: ...
 
-    def rpc_send(self, recipient_id: Hashable, msg: Any,) -> Any: ...
+    def rpc_send(
+        self,
+        recipient_id: Hashable,
+        msg: Any,
+    ) -> Any: ...
 
     def broadcast_send(
         self,

src/plumpy/message.py

@@ -7,7 +7,6 @@
 import logging
 from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Union, cast
 
-from plumpy.coordinator import Coordinator
 from plumpy.exceptions import PersistenceError, TaskRejectedError
 
 from . import loaders, persistence

src/plumpy/processes.py

@@ -34,9 +34,7 @@
     cast,
 )
 
-import kiwipy
-
-from plumpy.broadcast_filter import BroadcastFilter
+from plumpy.broadcast_filter import BroadcastFilter  # type: ignore
 from plumpy.coordinator import Coordinator
 
 try:
@@ -944,19 +942,16 @@ def _fire_event(self, evt: Callable[..., Any], *args: Any, **kwargs: Any) -> Non
     # region Communication
 
     def message_receive(self, msg: MessageType) -> Any:
-        """
-        Coroutine called when the process receives a message from the communicator
+        """Coroutine called when the process receives a message from the communicator
 
-        :param _comm: the communicator that sent the message
         :param msg: the message
         :return: the outcome of processing the message, the return value will be sent back as a response to the sender
         """
-        # self.logger.debug(
-        #     "Process<%s>: received RPC message with communicator '%s': %r",
-        #     self.pid,
-        #     _comm,
-        #     msg,
-        # )
+        self.logger.debug(
+            'Process<%s>: received RPC message: %r',
+            self.pid,
+            msg,
+        )
 
         intent = msg[message.INTENT_KEY]
 
@@ -977,19 +972,17 @@ def message_receive(self, msg: MessageType) -> Any:
     def broadcast_receive(
         self, msg: MessageType, sender: Any, subject: Any, correlation_id: Any
     ) -> Optional[concurrent.futures.Future]:
-        """
-        Coroutine called when the process receives a message from the communicator
+        """Coroutine called when the process receives a message from the communicator
 
         :param msg: the message
         """
+        self.logger.debug(
+            "Process<%s>: received broadcast message '%s': %r",
+            self.pid,
+            subject,
+            msg,
+        )
 
-        # self.logger.debug(
-        #     "Process<%s>: received broadcast message '%s' with communicator '%s': %r",
-        #     self.pid,
-        #     subject,
-        #     _comm,
-        #     msg,
-        # )
         # If we get a message we recognise then action it, otherwise ignore
         fn = None
         if subject == message.Intent.PLAY:

src/plumpy/rmq/communications.py

@@ -82,7 +82,9 @@ def converted(communicator: kiwipy.Communicator, *args: Any, **kwargs: Any) -> k
 CommT = TypeVar('CommT', bound=kiwipy.Communicator)
 
 
-def wrap_communicator(communicator: CommT, loop: Optional[asyncio.AbstractEventLoop] = None) -> 'LoopCommunicator[CommT]':
+def wrap_communicator(
+    communicator: CommT, loop: Optional[asyncio.AbstractEventLoop] = None
+) -> 'LoopCommunicator[CommT]':
     """
     Wrap a communicator such that all callbacks made to any subscribers are scheduled on the
     given event loop.

src/plumpy/rmq/process_control.py

@@ -9,6 +9,7 @@
 import kiwipy
 
 from plumpy import loaders
+from plumpy.controller import ProcessController
 from plumpy.coordinator import Coordinator
 from plumpy.message import (
     Intent,
@@ -29,8 +30,7 @@
 ProcessStatus = Any
 
 
-# FIXME: the class not fit typing of ProcessController protocol
-class RemoteProcessController:
+class RemoteProcessController(ProcessController):
     """
     Control remote processes using coroutines that will send messages and wait
     (in a non-blocking way) for their response
@@ -190,7 +190,7 @@ async def execute_process(
         return result
 
 
-class RemoteProcessThreadController:
+class RemoteProcessThreadController(ProcessController):
     """
     A class that can be used to control and launch remote processes
     """
@@ -212,7 +212,7 @@ def get_status(self, pid: 'PID_TYPE') -> kiwipy.Future:
         """
         return self._coordinator.rpc_send(pid, MessageBuilder.status())
 
-    def pause_process(self, pid: 'PID_TYPE', msg_text: str | None = None) -> kiwipy.Future:
+    def pause_process(self, pid: 'PID_TYPE', msg_text: str | None = None) -> Any:
         """Pause the process
 
         :param pid: the pid of the process to pause

tests/rmq/__init__.py

@@ -5,6 +5,7 @@
 import kiwipy
 import concurrent.futures
 
+from plumpy.coordinator import Coordinator
 from plumpy.exceptions import CoordinatorConnectionError
 
 if TYPE_CHECKING:
@@ -15,7 +16,7 @@
 
 
 @final
-class RmqCoordinator(Generic[U]):
+class RmqCoordinator(Coordinator, Generic[U]):
     def __init__(self, comm: U):
         self._comm = comm
 

tests/rmq/test_coordinator.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+from plumpy.coordinator import Coordinator
+from . import RmqCoordinator
+
+
+def test_mock_coordinator():
+    assert isinstance(RmqCoordinator, Coordinator)

tests/rmq/test_process_control.py

@@ -7,6 +7,7 @@
 from kiwipy import rmq
 
 import plumpy
+from plumpy.controller import ProcessController
 from plumpy.rmq import process_control
 
 from . import RmqCoordinator
@@ -42,6 +43,10 @@ def async_controller(_coordinator):
 def sync_controller(_coordinator):
     yield process_control.RemoteProcessThreadController(_coordinator)
 
+def test_remote_process_controller(sync_controller, async_controller):
+    assert isinstance(sync_controller, ProcessController)
+    assert isinstance(async_controller, ProcessController)
+
 
 class TestRemoteProcessController:
     @pytest.mark.asyncio

tests/test_coordinator.py

@@ -0,0 +1,7 @@
+# -*- coding: utf-8 -*-
+from plumpy.coordinator import Coordinator
+from .utils import MockCoordinator
+
+
+def test_mock_coordinator():
+    assert isinstance(MockCoordinator, Coordinator)

tests/utils.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 """Utilities for tests"""
+from __future__ import annotations
 
 import asyncio
 import collections
@@ -12,6 +13,7 @@
 
 import plumpy
 from plumpy import persistence, process_states, processes, utils
+from plumpy.coordinator import Coordinator
 from plumpy.exceptions import CoordinatorConnectionError
 from plumpy.message import MessageBuilder
 from plumpy.rmq import TaskRejected
@@ -24,7 +26,7 @@
 Snapshot = collections.namedtuple('Snapshot', ['state', 'bundle', 'outputs'])
 
 
-class MockCoordinator:
+class MockCoordinator(Coordinator):
     def __init__(self):
         self._task_receivers = {}
         self._broadcast_receivers = {}
@@ -105,7 +107,7 @@ def hook_broadcast_receiver(
         self._broadcast_receivers[identifier] = receiver
         return identifier
 
-    def unhook_broadcast_receiver(self, identifier: 'ID_TYPE | None') -> None: 
+    def unhook_broadcast_receiver(self, identifier: 'ID_TYPE | None') -> None:
         self._ensure_open()
         try:
             del self._broadcast_receivers[identifier]