diff --git a/monkey/monkey_island/cc/repository/i_event_repository.py b/monkey/monkey_island/cc/repository/i_event_repository.py index 9354afc59..94bc895dc 100644 --- a/monkey/monkey_island/cc/repository/i_event_repository.py +++ b/monkey/monkey_island/cc/repository/i_event_repository.py @@ -1,9 +1,11 @@ from abc import ABC, abstractmethod -from typing import Sequence, Type +from typing import Sequence, Type, TypeVar from common.events import AbstractAgentEvent from common.types import AgentID +T = TypeVar("T", bound=AbstractAgentEvent) + class IEventRepository(ABC): """A repository used to store and retrieve event objects""" @@ -27,9 +29,7 @@ class IEventRepository(ABC): """ @abstractmethod - def get_events_by_type( - self, event_type: Type[AbstractAgentEvent] - ) -> Sequence[AbstractAgentEvent]: + def get_events_by_type(self, event_type: Type[T]) -> Sequence[T]: """ Retrieve all events with same type