diff --git a/mistral/serialization.py b/mistral/serialization.py index db9351f01..6cd6330e8 100644 --- a/mistral/serialization.py +++ b/mistral/serialization.py @@ -137,6 +137,15 @@ class PolymorphicSerializer(Serializer): self.serializers[key] = serializer + def unregister(self, entity_cls): + key = self._get_serialization_key(entity_cls) + + if not key: + return + + if key in self.serializers: + del self.serializers[key] + def cleanup(self): self.serializers.clear() @@ -154,6 +163,11 @@ class PolymorphicSerializer(Serializer): serializer = self.serializers.get(key) + if not serializer: + raise RuntimeError( + "Failed to find a serializer for the key: %s" % key + ) + result = { '__serial_key': key, '__serial_data': serializer.serialize(entity) @@ -188,5 +202,9 @@ def register_serializer(entity_cls, serializer): get_polymorphic_serializer().register(entity_cls, serializer) +def unregister_serializer(entity_cls): + get_polymorphic_serializer().unregister(entity_cls) + + def cleanup(): get_polymorphic_serializer().cleanup() diff --git a/mistral/tests/unit/test_serialization.py b/mistral/tests/unit/test_serialization.py index 79e24f7c2..d66fee048 100644 --- a/mistral/tests/unit/test_serialization.py +++ b/mistral/tests/unit/test_serialization.py @@ -42,7 +42,7 @@ class SerializationTest(base.BaseTest): serialization.register_serializer(MyClass, MyClassSerializer()) - self.addCleanup(serialization.cleanup) + self.addCleanup(serialization.unregister_serializer, MyClass) def test_dict_based_serializer(self): obj = MyClass('a', 'b')