Added test of pop_traits_version

force_bdss/io/tests/test_workflow_writer.py

@@ -13,7 +13,8 @@ from force_bdss.io.workflow_reader import WorkflowReader
 from force_bdss.tests.dummy_classes.factory_registry_plugin import \
     DummyFactoryRegistryPlugin
 
-from force_bdss.io.workflow_writer import WorkflowWriter, traits_to_dict
+from force_bdss.io.workflow_writer import WorkflowWriter, traits_to_dict,\
+    pop_traits
 from force_bdss.core.workflow import Workflow
 
 
@@ -87,3 +88,14 @@ class TestWorkflowWriter(unittest.TestCase):
         mock_traits.__getstate__ = mock.Mock(return_value={"foo": "bar"})
 
         self.assertEqual(traits_to_dict(mock_traits), {"foo": "bar"})
+
+    def test_pop_traits_version(self):
+
+        test_dictionary = {'Entry1': {'Entry1-1': 4, '__traits_version__':67},
+                           'Entry2': [3, 'a', {'Entry2-1': 5,
+                                               '__traits_version__': 9001}],
+                           '__traits_version__': 13}
+        result_dictionary = {'Entry1': {'Entry1-1': 4, },
+                             'Entry2': [3, 'a', {'Entry2-1': 5, }], }
+        traitless_dictionary = pop_traits(test_dictionary)
+        self.assertEqual(traitless_dictionary,result_dictionary)

force_bdss/io/workflow_writer.py

@@ -93,29 +93,30 @@ def traits_to_dict(traits_obj):
     """Converts a traits class into a dict, removing the pesky
     traits version."""
 
-    def pop_traits(dictionary):
-        """Recursively remove the __traits_version__ attribute
-        from dictionary."""
-        try:
-            dictionary.pop("__traits_version__")
-        except KeyError:
-            pass
-
-        for key in dictionary:
-            # If we have a dict, remove the traits version
-            if isinstance(dictionary[key], dict):
-                pop_traits(dictionary[key])
-            # If we have a non-dict which contains a dict, remove traits from
-            # that as well
-            elif isinstance(dictionary[key], Iterable):
-                for element in dictionary[key]:
-                    if isinstance(element, dict):
-                        pop_traits(element)
-
-        return dictionary
-
     state = traits_obj.__getstate__()
 
-    state = pop_traits(state)
+    state = pop_traits_version(state)
 
     return state
+
+
+def pop_traits_version(dictionary):
+    """Recursively remove the __traits_version__ attribute
+    from dictionary."""
+    try:
+        dictionary.pop("__traits_version__")
+    except KeyError:
+        pass
+
+    for key in dictionary:
+        # If we have a dict, remove the traits version
+        if isinstance(dictionary[key], dict):
+            pop_traits(dictionary[key])
+        # If we have a non-dict which contains a dict, remove traits from
+        # that as well
+        elif isinstance(dictionary[key], Iterable):
+            for element in dictionary[key]:
+                if isinstance(element, dict):
+                    pop_traits(element)
+
+    return dictionary