Explicitly declare protocol implementations

Author: unkcpz

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
 

src/plumpy/coordinator.py

@@ -2,13 +2,14 @@
 from __future__ import annotations
 
 from re import Pattern
-from typing import TYPE_CHECKING, Any, Callable, Hashable, Protocol
+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,

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,9 +30,7 @@
 ProcessStatus = Any
 
 
-# This class not conform with typing of ProcessController protocol.
-# Does't matter too much, since this controller is not directly used as the controller by downstream.
-class RemoteProcessController:
+class RemoteProcessController(ProcessController):
     """
     Control remote processes using coroutines that will send messages and wait
     (in a non-blocking way) for their response
@@ -191,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
     """
@@ -213,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

@@ -13,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
@@ -25,7 +26,7 @@
 Snapshot = collections.namedtuple('Snapshot', ['state', 'bundle', 'outputs'])
 
 
-class MockCoordinator:
+class MockCoordinator(Coordinator):
     def __init__(self):
         self._task_receivers = {}
         self._broadcast_receivers = {}