From 1f899e9a84000c6475795a4e53e7581580177906 Mon Sep 17 00:00:00 2001 From: Steven Loria Date: Thu, 1 Nov 2018 20:41:28 -0400 Subject: [PATCH] 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/class_registry.py | 8 ++++++-- tests/test_registry.py | 10 +++++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/marshmallow/class_registry.py b/marshmallow/class_registry.py index 0d7abbced..1cf585d01 100644 --- 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 index c9d040b12..1f3c0be90 100644 --- a/tests/test_registry.py +++ b/tests/test_registry.py @@ -25,7 +25,7 @@ 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 @@ -33,7 +33,7 @@ def test_serializer_class_registry_register_same_classname_different_module(): #  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