diff --git a/force_bdss/io/workflow_writer.py b/force_bdss/io/workflow_writer.py index aadd19d2156dc98ca3ea2aff5db5fad7e5e40d0b..431343fbf028880d5d8f1b79294de260252c065d 100644 --- a/force_bdss/io/workflow_writer.py +++ b/force_bdss/io/workflow_writer.py @@ -1,5 +1,6 @@ import json from traits.api import HasStrictTraits +from collections import Iterable class WorkflowWriter(HasStrictTraits): @@ -91,10 +92,30 @@ class WorkflowWriter(HasStrictTraits): 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__() - try: - state.pop("__traits_version__") - except KeyError: - pass + + state = pop_traits(state) return state