diff --git a/monkey/common/event_queue/__init__.py b/monkey/common/event_queue/__init__.py index 04b67e625..3f1e149b6 100644 --- a/monkey/common/event_queue/__init__.py +++ b/monkey/common/event_queue/__init__.py @@ -1,4 +1,4 @@ -from .types import AgentEventSubscriber +from .types import AgentEventSubscriber, IslandEventSubscriber from .i_agent_event_queue import IAgentEventQueue from .i_island_event_queue import IIslandEventQueue, IslandEventTopic from .pypubsub_agent_event_queue import PyPubSubAgentEventQueue diff --git a/monkey/common/event_queue/i_island_event_queue.py b/monkey/common/event_queue/i_island_event_queue.py index a5614c37f..65eb94ee3 100644 --- a/monkey/common/event_queue/i_island_event_queue.py +++ b/monkey/common/event_queue/i_island_event_queue.py @@ -1,6 +1,8 @@ from abc import ABC, abstractmethod from enum import Enum -from typing import Any, Callable +from typing import Any + +from . import IslandEventSubscriber class IslandEventTopic(Enum): @@ -15,7 +17,7 @@ class IIslandEventQueue(ABC): """ @abstractmethod - def subscribe(self, topic: IslandEventTopic, subscriber: Callable[..., None]): + def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): """ Subscribes a subscriber to the specified event topic diff --git a/monkey/common/event_queue/pypubsub_island_event_queue.py b/monkey/common/event_queue/pypubsub_island_event_queue.py index 25be20633..d5d1df7cd 100644 --- a/monkey/common/event_queue/pypubsub_island_event_queue.py +++ b/monkey/common/event_queue/pypubsub_island_event_queue.py @@ -1,9 +1,9 @@ import logging -from typing import Any, Callable +from typing import Any from pubsub.core import Publisher -from . import IIslandEventQueue, IslandEventTopic +from . import IIslandEventQueue, IslandEventSubscriber, IslandEventTopic logger = logging.getLogger(__name__) @@ -13,7 +13,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): self._pypubsub_publisher = pypubsub_publisher self._refs = [] - def subscribe(self, topic: IslandEventTopic, subscriber: Callable[..., None]): + def subscribe(self, topic: IslandEventTopic, subscriber: IslandEventSubscriber): topic_value = topic.value # needs to be a string for pypubsub try: subscriber_name = subscriber.__name__ @@ -31,7 +31,7 @@ class PyPubSubIslandEventQueue(IIslandEventQueue): self._pypubsub_publisher.subscribe(topicName=topic_value, listener=subscriber) self._keep_subscriber_strongref(subscriber) - def _keep_subscriber_strongref(self, subscriber: Callable[..., None]): + def _keep_subscriber_strongref(self, subscriber: IslandEventSubscriber): # NOTE: PyPubSub stores subscribers by weak reference. From the documentation: # > PyPubSub holds listeners by weak reference so that the lifetime of the # > callable is not affected by PyPubSub: once the application no longer diff --git a/monkey/common/event_queue/types.py b/monkey/common/event_queue/types.py index 6b57ae712..bd0b23dff 100644 --- a/monkey/common/event_queue/types.py +++ b/monkey/common/event_queue/types.py @@ -3,3 +3,4 @@ from typing import Callable from common.events import AbstractAgentEvent AgentEventSubscriber = Callable[[AbstractAgentEvent], None] +IslandEventSubscriber = Callable[..., None]