From a32d9359b0722565d46a0efc99a770178897c4ad Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 18 Aug 2022 09:42:05 +0200 Subject: [PATCH] Common: Accept Union[str, Type[AbstractEvent]] in EventSerializerRegistry --- .../event_serializer_registry.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/monkey/common/event_serializers/event_serializer_registry.py b/monkey/common/event_serializers/event_serializer_registry.py index 2b49f4223..c5710c850 100644 --- a/monkey/common/event_serializers/event_serializer_registry.py +++ b/monkey/common/event_serializers/event_serializer_registry.py @@ -1,17 +1,21 @@ +from typing import Type, Union + from common.event_serializers import IEventSerializer +from common.events import AbstractEvent class EventSerializerRegistry: """ - Registry for event serializers + Registry for event serializers using event class. """ def __init__(self): self._registry = {} - def __setitem__(self, event_class_name: str, event_serializer: IEventSerializer): - self._registry[event_class_name] = event_serializer + def __setitem__( + self, event_class: Union[str, Type[AbstractEvent]], event_serializer: IEventSerializer + ): + self._registry[event_class] = event_serializer - def __getitem__(self, event_class_name: str): - event_serializer = self._registry[event_class_name] - return event_serializer + def __getitem__(self, event_class: Union[str, Type[AbstractEvent]]): + return self._registry[event_class]