diff --git a/force_bdss/io/tests/test_workflow_writer.py b/force_bdss/io/tests/test_workflow_writer.py index e973770dce1b50f35ce1100d7db9fc927e625c96..379af5a686bc03f3988ac56bf43aa857bfa02193 100644 --- a/force_bdss/io/tests/test_workflow_writer.py +++ b/force_bdss/io/tests/test_workflow_writer.py @@ -16,6 +16,7 @@ from force_bdss.tests.dummy_classes.factory_registry_plugin import \ from force_bdss.io.workflow_writer import WorkflowWriter, traits_to_dict,\ pop_recursive from force_bdss.core.workflow import Workflow +from force_bdss.core.input_slot_info import InputSlotInfo class TestWorkflowWriter(unittest.TestCase): @@ -89,7 +90,18 @@ class TestWorkflowWriter(unittest.TestCase): self.assertEqual(traits_to_dict(mock_traits), {"foo": "bar"}) - def test_pop_recursive(self): + def test_traits_to_dict(self): + + wfwriter = WorkflowWriter() + wf = self._create_workflow() + exec_layer = wf.execution_layers[0] + exec_layer.data_sources[0].input_slot_info = [InputSlotInfo()] + slotdata = exec_layer.data_sources[0].input_slot_info[0].__getstate__() + self.assertTrue("__traits_version__" in slotdata) + # Calls traits_to_dict for each data source + datastore_list = wfwriter._execution_layer_data(exec_layer) + new_slotdata = datastore_list[0]['model_data']['input_slot_info'] + self.assertTrue("__traits_version__" not in new_slotdata) test_dictionary = {'K1': {'K1': 'V1', 'K2': 'V2', 'K3': 'V3'}, 'K2': ['V1', 'V2', {'K1': 'V1', 'K2': 'V2', diff --git a/force_bdss/io/workflow_writer.py b/force_bdss/io/workflow_writer.py index a5e120e63b262a3260c3aebe45577b15210c2f5d..615793e1906a643e5d01a64bc2ea1db523746961 100644 --- a/force_bdss/io/workflow_writer.py +++ b/force_bdss/io/workflow_writer.py @@ -108,14 +108,14 @@ def pop_recursive(dictionary, remove_key): except KeyError: pass - for key in dictionary: + for key,value in dictionary.items(): # If remove_key is in the dict, remove it - if isinstance(dictionary[key], dict): - pop_recursive(dictionary[key], remove_key) + if isinstance(value, dict): + pop_recursive(value, remove_key) # If we have a non-dict iterable which contains a dict, # call pop.(remove_key) from that as well - elif isinstance(dictionary[key], Iterable): - for element in dictionary[key]: + elif isinstance(value, (tuple, list)): + for element in value: if isinstance(element, dict): pop_recursive(element, remove_key)