add support for child custom delegates

pymarshaler/__init__.py

@@ -1,4 +1,4 @@
-__version__ = '0.2.1'
+__version__ = '0.2.2'
 __all__ = ['Marshal', 'utils', 'arg_delegates', 'errors']
 
 from pymarshaler.marshal import Marshal

pymarshaler/marshal.py

@@ -18,16 +18,16 @@ def __init__(self):
         self.registered_delegates = {}
 
     def register(self, cls, delegate: ArgBuilderDelegate):
-        self.registered_delegates[cls.__name__] = delegate
+        self.registered_delegates[cls] = delegate
 
-    def get(self, cls):
-        return self.registered_delegates[cls.__name__]
-
-    def contains(self, cls):
+    def get_for(self, cls):
         try:
-            return cls.__name__ in self.registered_delegates
-        except AttributeError:
-            return False
+            for val in self.registered_delegates:
+                if cls == val or issubclass(cls, val):
+                    return self.registered_delegates[val]
+            return None
+        except TypeError:
+            return None
 
 
 class _ArgBuilderFactory:
@@ -53,8 +53,10 @@ def register(self, cls, delegate_cls):
         self._registered_delegates.register(cls, delegate_cls(cls))
 
     def get_delegate(self, cls) -> ArgBuilderDelegate:
-        if self._registered_delegates.contains(cls):
-            return self._registered_delegates.get(cls)
+        cls_maybe = self._registered_delegates.get_for(cls)
+
+        if cls_maybe:
+            return cls_maybe
         elif is_user_defined(cls):
             return self._default_arg_builder_delegates['UserDefined'](cls)
         elif '_name' in cls.__dict__:

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="pymarshaler",
-    version="0.2.1",
+    version="0.2.2",
     author="Hernan Romer",
     author_email="[email protected]",
     description="Package to marshal and unmarshal python objects",

tests/test_classes.py

@@ -90,6 +90,12 @@ def __init__(self):
         pass
 
 
+class ChildWithCustomDelegate(ClassWithCustomDelegate):
+
+    def __init__(self):
+        super().__init__()
+
+
 class CustomNoneDelegate(ArgBuilderDelegate):
 
     def __init__(self, cls):

tests/test_marshaling.py

@@ -122,6 +122,12 @@ def test_custom_delegate(self):
         result = marshal.unmarshal(ClassWithCustomDelegate, {})
         self.assertEqual(result, ClassWithCustomDelegate())
 
+    @timed
+    def test_child_custom_delegate(self):
+        marshal.register_delegate(ClassWithCustomDelegate, CustomNoneDelegate)
+        result = marshal.unmarshal(ChildWithCustomDelegate, {})
+        self.assertEqual(result, ChildWithCustomDelegate())
+
     @timed
     def test_nested_lists(self):
         nested_lists = NestedList([[Inner("Inner_1", 1)], [Inner("Inner_2", 2)]])