Move kiwipy/rmq related modules into a specific module

docs/source/nitpick-exceptions

@@ -23,7 +23,7 @@ py:class plumpy.base.state_machine.State
 py:class State
 py:class Process
 py:class plumpy.futures.CancellableAction
-py:class plumpy.communications.LoopCommunicator
+py:class plumpy.rmq.communications.LoopCommunicator
 py:class plumpy.persistence.PersistedPickle
 py:class plumpy.utils.AttributesFrozendict
 py:class plumpy.workchains._FunctionCall

src/plumpy/__init__.py

@@ -4,15 +4,14 @@
 
 import logging
 
-from .communications import *
 from .events import *
 from .exceptions import *
 from .futures import *
 from .loaders import *
+from .message import *
 from .mixins import *
 from .persistence import *
 from .ports import *
-from .process_comms import *
 from .process_listener import *
 from .process_states import *
 from .processes import *
@@ -27,8 +26,7 @@
     + futures.__all__
     + mixins.__all__
     + persistence.__all__
-    + communications.__all__
-    + process_comms.__all__
+    + message.__all__
     + process_listener.__all__
     + workchains.__all__
     + loaders.__all__

src/plumpy/coordinator.py

@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+from typing import Any, Callable, Pattern, Protocol
+
+RpcSubscriber = Callable[['Communicator', Any], Any]
+BroadcastSubscriber = Callable[['Communicator', Any, Any, Any, Any], Any]
+
+
+class Communicator(Protocol):
+    def add_rpc_subscriber(self, subscriber: RpcSubscriber, identifier=None) -> Any: ...
+
+    def add_broadcast_subscriber(
+        self, subscriber: BroadcastSubscriber, subject_filter: str | Pattern[str] | None = None, identifier=None
+    ) -> Any: ...
+
+    def remove_rpc_subscriber(self, identifier): ...
+
+    def remove_broadcast_subscriber(self, identifier): ...
+
+    def broadcast_send(self, body, sender=None, subject=None, correlation_id=None) -> bool: ...

src/plumpy/exceptions.py

@@ -1,16 +1,21 @@
 # -*- coding: utf-8 -*-
 from typing import Optional
 
-__all__ = ['ClosedError', 'InvalidStateError', 'KilledError', 'PersistenceError', 'UnsuccessfulResult']
+__all__ = [
+    'ClosedError',
+    'InvalidStateError',
+    'KilledError',
+    'PersistenceError',
+    'UnsuccessfulResult',
+]
 
 
 class KilledError(Exception):
     """The process was killed."""
 
 
 class InvalidStateError(Exception):
-    """
-    Raised when an operation is attempted that requires the process to be in a state
+    """Raised when an operation is attempted that requires the process to be in a state
     that is different from the current state
     """
 
@@ -33,3 +38,7 @@ class PersistenceError(Exception):
 
 class ClosedError(Exception):
     """Raised when an mutable operation is attempted on a closed process"""
+
+
+class TaskRejectedError(Exception):
+    """A task was rejected by the coordinacor"""

src/plumpy/futures.py

@@ -4,24 +4,33 @@
 """
 
 import asyncio
-from typing import Any, Awaitable, Callable, Optional
+import contextlib
+from typing import Any, Awaitable, Callable, Generator, Optional
 
-import kiwipy
+__all__ = ['CancellableAction', 'capture_exceptions', 'create_task', 'create_task']
 
-__all__ = ['CancelledError', 'Future', 'chain', 'copy_future', 'create_task', 'gather']
 
-CancelledError = kiwipy.CancelledError
+class InvalidFutureError(Exception):
+    """Exception for when a future or action is in an invalid state"""
 
 
-class InvalidStateError(Exception):
-    """Exception for when a future or action is in an invalid state"""
+Future = asyncio.Future
 
 
-copy_future = kiwipy.copy_future
-chain = kiwipy.chain
-gather = asyncio.gather
+@contextlib.contextmanager
+def capture_exceptions(future: Future[Any], ignore: tuple[type[BaseException], ...] = ()) -> Generator[None, Any, None]:
+    """
+    Capture any exceptions in the context and set them as the result of the given future
 
-Future = asyncio.Future
+    :param future: The future to the exception on
+    :param ignore: An optional list of exception types to ignore, these will be raised and not set on the future
+    """
+    try:
+        yield
+    except ignore:
+        raise
+    except Exception as exception:
+        future.set_exception(exception)
 
 
 class CancellableAction(Future):
@@ -46,10 +55,10 @@ def run(self, *args: Any, **kwargs: Any) -> None:
         :param kwargs: the keyword arguments to the action
         """
         if self.done():
-            raise InvalidStateError('Action has already been ran')
+            raise InvalidFutureError('Action has already been ran')
 
         try:
-            with kiwipy.capture_exceptions(self):
+            with capture_exceptions(self):
                 self.set_result(self._action(*args, **kwargs))
         finally:
             self._action = None  # type: ignore
@@ -70,38 +79,9 @@ def create_task(coro: Callable[[], Awaitable[Any]], loop: Optional[asyncio.Abstr
     future = loop.create_future()
 
     async def run_task() -> None:
-        with kiwipy.capture_exceptions(future):
+        with capture_exceptions(future):
             res = await coro()
             future.set_result(res)
 
     asyncio.run_coroutine_threadsafe(run_task(), loop)
     return future
-
-
-def unwrap_kiwi_future(future: kiwipy.Future) -> kiwipy.Future:
-    """
-    Create a kiwi future that represents the final results of a nested series of futures,
-    meaning that if the futures provided itself resolves to a future the returned
-    future will not resolve to a value until the final chain of futures is not a future
-    but a concrete value.  If at any point in the chain a future resolves to an exception
-    then the returned future will also resolve to that exception.
-
-    :param future: the future to unwrap
-    :return: the unwrapping future
-
-    """
-    unwrapping = kiwipy.Future()
-
-    def unwrap(fut: kiwipy.Future) -> None:
-        if fut.cancelled():
-            unwrapping.cancel()
-        else:
-            with kiwipy.capture_exceptions(unwrapping):
-                result = fut.result()
-                if isinstance(result, kiwipy.Future):
-                    result.add_done_callback(unwrap)
-                else:
-                    unwrapping.set_result(result)
-
-    future.add_done_callback(unwrap)
-    return unwrapping