Fix memory leak when dynamically creating schema classes

* Don't add duplicate entries for the same full path * Fix behavior when classes have the same name but different module close #732
marshmallow-code · Nov 2, 2018 · 1f899e9 · 1f899e9
1 parent 510fd81
commit 1f899e9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 7 deletions.
diff --git a/marshmallow/class_registry.py b/marshmallow/class_registry.py
@@ -48,11 +48,15 @@ class MyClass:
     if classname in _registry and not \
             any(each.__module__ == module for each in _registry[classname]):
         _registry[classname].append(cls)
-    else:
+    elif classname not in _registry:
         _registry[classname] = [cls]
 
     # Also register the full path
-    _registry.setdefault(fullpath, []).append(cls)
+    if fullpath not in _registry:
+        _registry.setdefault(fullpath, []).append(cls)
+    else:
+        # If fullpath does exist, replace existing entry
+        _registry[fullpath] = [cls]
     return None
 
 def get_class(classname, all=False):

diff --git a/tests/test_registry.py b/tests/test_registry.py
@@ -25,15 +25,15 @@ def test_serializer_class_registry_register_same_classname_different_module():
 
     reglen = len(class_registry._registry)
 
-    cls1 = type('MyTestRegSchema', (Schema,), {'__module__': 'modA'})
+    type('MyTestRegSchema', (Schema,), {'__module__': 'modA'})
 
     assert 'MyTestRegSchema' in class_registry._registry
     assert len(class_registry._registry.get('MyTestRegSchema')) == 1
     assert 'modA.MyTestRegSchema' in class_registry._registry
     #  storing for classname and fullpath
     assert len(class_registry._registry) == reglen + 2
 
-    cls2 = type('MyTestRegSchema', (Schema,), {'__module__': 'modB'})
+    type('MyTestRegSchema', (Schema,), {'__module__': 'modB'})
 
     assert 'MyTestRegSchema' in class_registry._registry
     #  aggregating classes with same name from different modules
@@ -42,7 +42,7 @@ def test_serializer_class_registry_register_same_classname_different_module():
     #  storing for same classname (+0) and different module (+1)
     assert len(class_registry._registry) == reglen + 2 + 1
 
-    cls2 = type('MyTestRegSchema', (Schema,), {'__module__': 'modB'})
+    type('MyTestRegSchema', (Schema,), {'__module__': 'modB'})
 
     assert 'MyTestRegSchema' in class_registry._registry
     #  only the class with matching module has been replaced
@@ -56,7 +56,7 @@ def test_serializer_class_registry_override_if_same_classname_same_module():
 
     reglen = len(class_registry._registry)
 
-    cls1 = type('MyTestReg2Schema', (Schema,), {'__module__': 'SameModulePath'})
+    type('MyTestReg2Schema', (Schema,), {'__module__': 'SameModulePath'})
 
     assert 'MyTestReg2Schema' in class_registry._registry
     assert len(class_registry._registry.get('MyTestReg2Schema')) == 1
@@ -65,7 +65,7 @@ def test_serializer_class_registry_override_if_same_classname_same_module():
     #  storing for classname and fullpath
     assert len(class_registry._registry) == reglen + 2
 
-    cls2 = type('MyTestReg2Schema', (Schema,), {'__module__': 'SameModulePath'})
+    type('MyTestReg2Schema', (Schema,), {'__module__': 'SameModulePath'})
 
     assert 'MyTestReg2Schema' in class_registry._registry
     #  overriding same class name and same module