From b541dc465d55d5b19f5029f57312faa5466f8d14 Mon Sep 17 00:00:00 2001 From: Ilija Lazoroski Date: Thu, 18 Aug 2022 15:21:33 +0200 Subject: [PATCH] Common: Type checking in EventSerializerRegistry --- .../event_serializer_registry.py | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/monkey/common/event_serializers/event_serializer_registry.py b/monkey/common/event_serializers/event_serializer_registry.py index 3e1ada7e8..fd02045d1 100644 --- a/monkey/common/event_serializers/event_serializer_registry.py +++ b/monkey/common/event_serializers/event_serializer_registry.py @@ -7,13 +7,35 @@ from common.events import AbstractEvent class EventSerializerRegistry: """ Registry for event serializers using event class. + + Example: + event_serializer_registry = EventSerializerRegistry() + event_serializer_registry[MyEvent] = MyEventSerializer() + + my_event_dict = {"type": "MyEvent", "data": "123"} + + serializer = event_serializer_registry[my_event_dict["type"]] + my_event_object = serializer.deserialize(my_event_dict) """ def __init__(self): self._registry = {} def __setitem__(self, event_class: Type[AbstractEvent], event_serializer: IEventSerializer): - self._registry[event_class] = event_serializer + if not issubclass(event_class, AbstractEvent): + raise TypeError(f"Event class must be of type: {AbstractEvent.__name__}") + + if not isinstance(event_serializer, IEventSerializer): + raise TypeError(f"Event serializer must be of type: {IEventSerializer.__name__}") + + self._registry[event_class] = event_serializer + self._registry[event_class.__name__] = event_serializer + + def __getitem__(self, event_class: Union[str, Type[AbstractEvent]]) -> IEventSerializer: + if not (isinstance(event_class, str) or issubclass(event_class, AbstractEvent)): + raise TypeError( + f"Registry get key {event_class} must be of type: {AbstractEvent.__name__} or " + f"{str.__name__}" + ) - def __getitem__(self, event_class: Union[str, Type[AbstractEvent]]): return self._registry[event_class]