From b6c026efd266beecbabfa80a0f18991cd5fc416e Mon Sep 17 00:00:00 2001 From: Spyros Zoupanos Date: Sun, 24 Mar 2019 11:45:54 +0000 Subject: [PATCH] There were problems when overriding the default behaviour of the rules. Moreover they were moved to separate methods to resuce the complexity of the export method. Various tests were added: - Export tests that are based on a real workchain were added to ensure that a complex and real AiiDA graph based on workchains is exported correctly. - Export tests based on a synthetic graph were added that explicitly check every flag that changes the default behaviour of the export set expansion rules. --- aiida/backends/tests/export_and_import.py | 505 +++++++++---- aiida/common/graph.py | 13 +- aiida/orm/importexport.py | 846 +++++++++++----------- docs/source/nitpick-exceptions | 1 + 4 files changed, 833 insertions(+), 532 deletions(-) diff --git a/aiida/backends/tests/export_and_import.py b/aiida/backends/tests/export_and_import.py index 0738000881..8f3cf12e2f 100644 --- a/aiida/backends/tests/export_and_import.py +++ b/aiida/backends/tests/export_and_import.py @@ -13,6 +13,8 @@ from aiida.backends.testbase import AiidaTestCase from aiida.orm.importexport import import_data +from aiida.orm.calculation.inline import make_inline + class TestSpecificImport(AiidaTestCase): @@ -169,6 +171,7 @@ def test_cycle_structure_data(self): qb.append(Calculation, output_of='remote') self.assertGreater(len(qb.all()), 0) + class TestSimple(AiidaTestCase): def setUp(self): @@ -746,11 +749,9 @@ def test_workcalculation_2(self): from aiida.common.links import LinkType from aiida.orm.importexport import export - from aiida.common.exceptions import NotExistent # Creating a folder for the import/export files temp_folder = tempfile.mkdtemp() - try: master = WorkCalculation().store() slave = WorkCalculation().store() @@ -778,7 +779,6 @@ def test_workcalculation_2(self): # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True) - def test_reexport(self): """ Export something, import and reexport and check if everything is valid. @@ -896,6 +896,7 @@ def get_hash_from_db_content(groupname): # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True) + class TestComplex(AiidaTestCase): def test_complex_graph_import_export(self): @@ -981,6 +982,7 @@ def test_complex_graph_import_export(self): # Deleting the created temporary folder shutil.rmtree(temp_folder, ignore_errors=True) + class TestComputer(AiidaTestCase): def setUp(self): @@ -1461,6 +1463,7 @@ def test_import_of_django_sqla_export_file(self): comp1_metadata, "Not the expected metadata were found") + class TestLinks(AiidaTestCase): def setUp(self): @@ -1481,6 +1484,176 @@ def get_all_node_links(self): edge_project=['label', 'type'], output_of='input') return qb.all() + class ComplexGraph: + """ + A class responsible to create a "complex" graph with all available link types + (INPUT, CREATE, RETURN and CALL) + """ + # Information about the size of the available scenarios (dict) + export_scenarios_info = None + + # Information regarding the export flags that are overridden for each scenario + export_scenarios_flags = None + + def __init__(self): + # Information for the size of sub scenarios of every export scenario + self.export_scenarios_info = dict() + self.export_scenarios_info['default'] = 9 + self.export_scenarios_info['input_forward_true'] = 3 + self.export_scenarios_info['create_reversed_false'] = 4 + self.export_scenarios_info['return_reversed_true'] = 4 + self.export_scenarios_info['call_reversed_true'] = 2 + + # The export flags of every export scenario + self.export_scenarios_flags = dict() + self.export_scenarios_flags['default'] = dict() + self.export_scenarios_flags['input_forward_true'] = {'input_forward': True} + self.export_scenarios_flags['create_reversed_false'] = {'create_reversed': False} + self.export_scenarios_flags['return_reversed_true'] = {'return_reversed': True} + self.export_scenarios_flags['call_reversed_true'] = {'call_reversed': True} + + def construct_complex_graph(self, computer, scenario_name='default', sub_scenario_number=0): + """ + This method creates the graph and returns its nodes. + It also returns various combinations of nodes that need to be extracted + but also the final expected set of nodes (after adding the expected + predecessors, successors etc) based on a set of arguments. + :param computer: The computer that will be used for the calculations of the + generated graph + :param scenario_name: Various scenarios that we can test. These correspond + to the set expansion rules that we would like to override. The available + options are: default, input_forward_true, create_reversed_false, + return_reversed_true and call_reversed_true. + :param sub_scenario_number: Based on the main scenario that we have chosen + the sub-scenario (test case) that we would like to get for testing. + :return: The graph nodes and the sub-scenario that corresponds to the provided + arguments. The sub-scenario is composes of a set of export nodes and a set + of nodes that should have been finally exported after applying the set + expansion rules. + """ + from aiida.orm.data.base import Int + from aiida.orm.calculation.job import JobCalculation + from aiida.orm.calculation.work import WorkCalculation + from aiida.common.datastructures import calc_states + from aiida.common.links import LinkType + + # Node creation + d1 = Int(1).store() + d2 = Int(2).store() + wc1 = WorkCalculation().store() + wc2 = WorkCalculation().store() + + pw1 = JobCalculation() + pw1.set_computer(computer) + pw1.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) + pw1.store() + + d3 = Int(3).store() + d4 = Int(4).store() + + pw2 = JobCalculation() + pw2.set_computer(computer) + pw2.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) + pw2.store() + + d5 = Int(5).store() + d6 = Int(6).store() + + # Link creation + wc1.add_link_from(d1, 'input1', link_type=LinkType.INPUT) + wc1.add_link_from(d2, 'input2', link_type=LinkType.INPUT) + + wc2.add_link_from(d1, 'input', link_type=LinkType.INPUT) + wc2.add_link_from(wc1, 'call', link_type=LinkType.CALL) + + pw1.add_link_from(d1, 'input', link_type=LinkType.INPUT) + pw1.add_link_from(wc2, 'call', link_type=LinkType.CALL) + pw1._set_state(calc_states.PARSING) + + d3.add_link_from(pw1, 'create', link_type=LinkType.CREATE) + d3.add_link_from(wc2, 'return', link_type=LinkType.RETURN) + + d4.add_link_from(pw1, 'create', link_type=LinkType.CREATE) + d4.add_link_from(wc2, 'return', link_type=LinkType.RETURN) + + pw2.add_link_from(d4, 'input', link_type=LinkType.INPUT) + pw2._set_state(calc_states.PARSING) + + d5.add_link_from(pw2, 'create', link_type=LinkType.CREATE) + d6.add_link_from(pw2, 'create', link_type=LinkType.CREATE) + + # Return the generated nodes + graph_nodes = [d1, d2, d3, d4, d5, d6, pw1, pw2, wc1, wc2] + + # Create various combinations of nodes that should be exported + # and the final set of nodes that are exported in each case, following + # predecessor/successor links. + + # The export list for default values on the export flags + export_list_default = [ + (wc1, [d1, d2, d3, d4, pw1, wc1, wc2]), + (wc2, [d1, d3, d4, pw1, wc2]), + (d3, [d1, d3, d4, pw1]), + (d4, [d1, d3, d4, pw1]), + (d5, [d1, d3, d4, d5, d6, pw1, pw2]), + (pw1, [d1, d3, d4, pw1]), + (pw2, [d1, d3, d4, d5, d6, pw1, pw2]), + (d1, [d1]), + (d2, [d2]), + ] + + # The export list of selected nodes for the export flags: input_forward = True + export_list_input_forward_true = [ + (d1, [d1, d2, d3, d4, d5, d6, wc1, wc2, pw1, pw2]), + (d2, [d1, d2, d3, d4, d5, d6, wc1, wc2, pw1, pw2]), + (d4, [d1, d2, d3, d4, d5, d6, wc1, wc2, pw1, pw2]), + ] + + # The export list of selected nodes for the export flags: create_reversed = False + export_list_create_reversed_false = [ + (d3, [d3]), + (d4, [d4]), + (d5, [d5]), + (d6, [d6]), + ] + + # The export list of selected nodes for the export flags: return_reversed = True + export_list_return_reversed_true = [ + (d3, [d1, d3, d4, pw1, wc2]), + (d4, [d1, d3, d4, pw1, wc2]), + (d5, [d1, d3, d4, d5, d6, pw1, pw2, wc2]), + (d6, [d1, d3, d4, d5, d6, pw1, pw2, wc2]), + ] + + # The export list of selected nodes for the export flags: call_reversed = True + export_list_call_reversed_true = [ + (pw1, [d1, d2, d3, d4, wc1, wc2, pw1]), + (wc2, [d1, d2, d3, d4, wc1, wc2, pw1]), + ] + + export_scenarios = dict() + export_scenarios['default'] = export_list_default + export_scenarios['input_forward_true'] = export_list_input_forward_true + export_scenarios['create_reversed_false'] = export_list_create_reversed_false + export_scenarios['return_reversed_true'] = export_list_return_reversed_true + export_scenarios['call_reversed_true'] = export_list_call_reversed_true + + return graph_nodes, export_scenarios[scenario_name][sub_scenario_number] + + def get_scenarios_names_and_size(self): + """ + The available scenarios and the number of the sub-scenarios + """ + return self.export_scenarios_info + + def get_scenarios_export_flags(self, scenario_name): + """ + The flags that correspond to the given scenario. These flags should be passed + to the export method/function to override the default behaviour of the set + expansion rules. + """ + return self.export_scenarios_flags[scenario_name] + def test_input_and_create_links(self): """ Simple test that will verify that INPUT and CREATE links are properly exported and @@ -1520,88 +1693,6 @@ def test_input_and_create_links(self): finally: shutil.rmtree(tmp_folder, ignore_errors=True) - def construct_complex_graph(self, export_combination = 0): - """ - This method creates a "complex" graph with all available link types - (INPUT, CREATE, RETURN and CALL) and returns the nodes of the graph. It - also returns various combinations of nodes that need to be extracted - but also the final expected set of nodes (after adding the expected - predecessors, desuccessors). - """ - from aiida.orm.data.base import Int - from aiida.orm.calculation.job import JobCalculation - from aiida.orm.calculation.work import WorkCalculation - from aiida.common.datastructures import calc_states - from aiida.common.links import LinkType - - if export_combination < 0 or export_combination > 8: - return None - - # Node creation - d1 = Int(1).store() - d2 = Int(1).store() - wc1 = WorkCalculation().store() - wc2 = WorkCalculation().store() - - pw1 = JobCalculation() - pw1.set_computer(self.computer) - pw1.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) - pw1.store() - - d3 = Int(1).store() - d4 = Int(1).store() - - pw2 = JobCalculation() - pw2.set_computer(self.computer) - pw2.set_resources({"num_machines": 1, "num_mpiprocs_per_machine": 1}) - pw2.store() - - d5 = Int(1).store() - d6 = Int(1).store() - - # Link creation - wc1.add_link_from(d1, 'input1', link_type=LinkType.INPUT) - wc1.add_link_from(d2, 'input2', link_type=LinkType.INPUT) - - wc2.add_link_from(d1, 'input', link_type=LinkType.INPUT) - wc2.add_link_from(wc1, 'call', link_type=LinkType.CALL) - - pw1.add_link_from(d1, 'input', link_type=LinkType.INPUT) - pw1.add_link_from(wc2, 'call', link_type=LinkType.CALL) - pw1._set_state(calc_states.PARSING) - - d3.add_link_from(pw1, 'create', link_type=LinkType.CREATE) - d3.add_link_from(wc2, 'return', link_type=LinkType.RETURN) - - d4.add_link_from(pw1, 'create', link_type=LinkType.CREATE) - d4.add_link_from(wc2, 'return', link_type=LinkType.RETURN) - - pw2.add_link_from(d4, 'input', link_type=LinkType.INPUT) - pw2._set_state(calc_states.PARSING) - - d5.add_link_from(pw2, 'create', link_type=LinkType.CREATE) - d6.add_link_from(pw2, 'create', link_type=LinkType.CREATE) - - # Return the generated nodes - graph_nodes = [d1, d2, d3, d4, d5, d6, pw1, pw2, wc1, wc2] - - # Create various combinations of nodes that should be exported - # and the final set of nodes that are exported in each case, following - # predecessor/successor links. - export_list = [ - (wc1, [d1, d2, d3, d4, pw1, wc1, wc2]), - (wc2, [d1, d3, d4, pw1, wc2]), - (d3, [d1, d3, d4, pw1]), - (d4, [d1, d3, d4, pw1]), - (d5, [d1, d3, d4, d5, d6, pw1, pw2]), - (d6, [d1, d3, d4, d5, d6, pw1, pw2]), - (pw2, [d1, d3, d4, d5, d6, pw1, pw2]), - (d1, [d1]), - (d2, [d2]) - ] - - return graph_nodes, export_list[export_combination] - def test_complex_workflow_graph_links(self): """ This test checks that all the needed links are correctly exported and @@ -1618,7 +1709,7 @@ def test_complex_workflow_graph_links(self): tmp_folder = tempfile.mkdtemp() try: - graph_nodes, _ = self.construct_complex_graph() + graph_nodes, _ = self.ComplexGraph().construct_complex_graph(self.computer) # Getting the input, create, return and call links qb = QueryBuilder() @@ -1648,49 +1739,55 @@ def test_complex_workflow_graph_links(self): shutil.rmtree(tmp_folder, ignore_errors=True) def test_complex_workflow_graph_export_set_expansion(self): + """ + Test the various export set_expansion rules. It tests the default behaviour but also + all the flags that oveeride this behaviour on a manually created AiiDA graph. + """ import os, shutil, tempfile from aiida.orm.importexport import export from aiida.orm.querybuilder import QueryBuilder from aiida.orm import Node - for export_conf in range(0,8): - - graph_nodes, (export_node, export_target) = ( - self.construct_complex_graph(export_conf)) - - tmp_folder = tempfile.mkdtemp() - try: - export_file = os.path.join(tmp_folder, 'export.tar.gz') - export([export_node], outfile=export_file, silent=True) - export_node_str = str(export_node) - - self.clean_db() - self.insert_data() - - import_data(export_file, silent=True) - - # Get all the nodes of the database - qb = QueryBuilder() - qb.append(Node, project='uuid') - imported_node_uuids = set(str(_[0]) for _ in qb.all()) - - export_target_uuids = set(str(_.uuid) for _ in export_target) - - from aiida.orm.utils import load_node - self.assertEquals( - export_target_uuids, - imported_node_uuids, - "Problem in comparison of export node: " + - str(export_node_str) + "\n" + - "Expected set: " + str(export_target_uuids) + "\n" + - "Imported set: " + str(imported_node_uuids) + "\n" + - "Difference: " + str([load_node(_) for _ in - export_target_uuids.symmetric_difference( - imported_node_uuids)]) - ) + cg = self.ComplexGraph() + for scenario_name, scenarios_number in cg.get_scenarios_names_and_size().iteritems(): + for sub_scenarion_no in range(0, scenarios_number): + _, (export_node, export_target) = cg.construct_complex_graph( + self.computer, scenario_name=scenario_name, sub_scenario_number=sub_scenarion_no) + tmp_folder = tempfile.mkdtemp() - finally: - shutil.rmtree(tmp_folder, ignore_errors=True) + try: + export_file = os.path.join(tmp_folder, 'export.tar.gz') + export([export_node], outfile=export_file, + silent=True, **cg.get_scenarios_export_flags(scenario_name)) + export_node_str = str(export_node) + export_node_type = str(export_node.type) + + self.clean_db() + self.insert_data() + + import_data(export_file, silent=True) + + # Get all the nodes of the database + qb = QueryBuilder() + qb.append(Node, project='uuid') + imported_node_uuids = set(str(_[0]) for _ in qb.all()) + + export_target_uuids = set(str(_.uuid) for _ in export_target) + + from aiida.orm.utils import load_node + self.assertEquals( + export_target_uuids, + imported_node_uuids, + "Problem in comparison of export node (Export flags: {}): ".format(scenario_name) + + str(export_node_str) + ". Export node type:" + export_node_type + "\n" + + "Expected set: " + str(export_target_uuids) + "\n" + + "Imported set: " + str(imported_node_uuids) + "\n" + + "Difference: " + str(export_target_uuids.symmetric_difference( + imported_node_uuids)) + ) + + finally: + shutil.rmtree(tmp_folder, ignore_errors=True) def test_recursive_export_input_and_create_links_proper(self): """ @@ -1764,7 +1861,6 @@ def test_recursive_export_input_and_create_links_proper(self): import_data(export_file, silent=True) import_links = self.get_all_node_links() - export_set = [tuple(_) for _ in export_links] import_set = [tuple(_) for _ in import_links] @@ -1980,4 +2076,163 @@ def test_that_input_code_is_exported_correctly(self): "code to the calculation node. {} found." .format(qb.count())) finally: - shutil.rmtree(tmp_folder, ignore_errors=True) \ No newline at end of file + shutil.rmtree(tmp_folder, ignore_errors=True) + + +@make_inline +def sum_inline(x, y): + from aiida.orm.data.base import Int + return {'result': Int(x + y)} + + +class TestRealWorkChainExport(AiidaTestCase): + """ + Create an AiiDA graph by executing a workchain with sub-workchains and test if all the graph is + exported by setting all the set expansion rules flags to True. + """ + from aiida.work.workchain import WorkChain + + def check_correct_export(self, uuids_of_target_export_nodes, uuids_of_target_import_nodes, + starting_node_to_draw_graph=None, **export_args): + import os, shutil, tempfile + from aiida.orm.importexport import export + from aiida.orm.querybuilder import QueryBuilder + from aiida.orm.node import Node + from aiida.orm import load_node + from aiida.common.graph import draw_graph + + try: + tmp_folder = tempfile.mkdtemp() + + target_export_nodes = set() + for uuid in uuids_of_target_export_nodes: + target_export_nodes.add(load_node(uuid)) + + if starting_node_to_draw_graph is not None: + draw_graph(load_node(starting_node_to_draw_graph), filename_suffix='before_export', format='pdf') + + export_file = os.path.join(tmp_folder, 'export.tar.gz') + export(target_export_nodes, outfile=export_file, silent=True, **export_args) + + self.clean_db() + self.insert_data() + + import_data(export_file, silent=True, ignore_unknown_nodes=True) + + if starting_node_to_draw_graph is not None: + draw_graph(load_node(starting_node_to_draw_graph), filename_suffix='after_import', format='pdf') + + # Check that the expected nodes are imported + uuids_in_db = set([str(uuid) for [uuid] in QueryBuilder().append(Node, project=['uuid']).all()]) + self.assertTrue(uuids_of_target_import_nodes.issubset(uuids_in_db), "The expected nodes ({}) where not " + "found at the imported graph ({})." + .format(uuids_of_target_import_nodes, uuids_in_db)) + finally: + shutil.rmtree(tmp_folder, ignore_errors=True) + + class SubWorkChain(WorkChain): + """ + A Workchain that calls itself for various levels. + """ + @classmethod + def define(cls, spec): + from aiida.orm.data.base import Int + + super(TestRealWorkChainExport.SubWorkChain, cls).define(spec) + spec.input("wf_counter", valid_type=Int) + spec.input('a', valid_type=Int) + spec.input('b', valid_type=Int) + spec.outline( + cls.start + ) + spec.output('result', valid_type=Int) + + def start(self): + from aiida.work.run import run + from aiida.orm.data.base import Int + + if self.inputs.wf_counter > 0: + inputs = { + 'wf_counter': Int(self.inputs.wf_counter - 1), + 'a': self.inputs.a, + 'b': self.inputs.b + } + result = run(TestRealWorkChainExport.SubWorkChain, **inputs) + self.out('result', result['result']) + else: + node, result = sum_inline(x=self.inputs.a, y=self.inputs.b) + self.out('result', result['result']) + + def setUp(self): + import aiida.work.util as util + import plum + super(TestRealWorkChainExport, self).setUp() + self.assertEquals(len(util.ProcessStack.stack()), 0) + self.assertEquals(len(plum.process_monitor.MONITOR.get_pids()), 0) + + def tearDown(self): + import aiida.work.util as util + import plum + super(TestRealWorkChainExport, self).tearDown() + self.assertEquals(len(util.ProcessStack.stack()), 0) + self.assertEquals(len(plum.process_monitor.MONITOR.get_pids()), 0) + self.clean_db() + self.insert_data() + + def test_any_wc(self): + """ + Check that all the WorkCalculations can be exported when I start the export from every WorkCalculation + (with the flags call_reversed=True, return_reversed=True, create_reverese=True) + """ + from aiida.orm.data.base import Int + from aiida.orm.querybuilder import QueryBuilder + from aiida.orm.implementation.general.calculation.work import WorkCalculation + from aiida.work.run import run + + _, _ = run(TestRealWorkChainExport.SubWorkChain, a=Int(1), b=Int(2), + wf_counter=Int(2), _return_pid=True) + # All the WorkCalculation UUIDs + tagret_wc_uuids = set([str(uuid) for [uuid] in QueryBuilder().append(WorkCalculation, project=['uuid']).all()]) + # For every WorkCalculation, I should export (and import correctly) all the WorkCalculations + for wf_uuid in tagret_wc_uuids: + self.check_correct_export([wf_uuid], tagret_wc_uuids, call_reversed=True, return_reversed=True, + create_reversed=True) + + def test_top_wc(self): + """ + Check that all the WorkCalculations are exported when the top WorkCalculation is used to start the export + (without providing any extra args) + """ + from aiida.orm.data.base import Int + from aiida.orm.querybuilder import QueryBuilder + from aiida.orm.implementation.general.calculation.work import WorkCalculation + from aiida.work.run import run + from aiida.orm import load_node + + ret_node, top_wf_pk = run(TestRealWorkChainExport.SubWorkChain, a=Int(1), b=Int(2), + wf_counter=Int(2), _return_pid=True) + # All the WorkCalculation UUIDs + tagret_wc_uuids = set([str(uuid) for [uuid] in QueryBuilder().append(WorkCalculation, project=['uuid']).all()]) + + # Check that if we export the top WorkCalculation, all the graph is exported + self.check_correct_export([load_node(top_wf_pk).uuid], tagret_wc_uuids) + + def test_lower_return_node(self): + """ + Check that if we export the lower return node all the WorkCalculations are exported + (with the flags call_reversed=True, return_reversed=True, create_reveresed=True) + """ + from aiida.orm.data.base import Int + from aiida.orm.querybuilder import QueryBuilder + from aiida.orm.implementation.general.calculation.work import WorkCalculation + from aiida.work.run import run + from aiida.orm.node import Node + + ret_node, _ = run(TestRealWorkChainExport.SubWorkChain, a=Int(1), b=Int(2), + wf_counter=Int(2), _return_pid=True) + + # All the WorkCalculation UUIDs + tagret_wc_uuids = set([str(uuid) for [uuid] in QueryBuilder().append(WorkCalculation, project=['uuid']).all()]) + # Check that if we export the lower return node, all the graph is exported + self.check_correct_export([ret_node['result'].uuid], tagret_wc_uuids, call_reversed=True, return_reversed=True, + create_reversed=True, input_forward=True) diff --git a/aiida/common/graph.py b/aiida/common/graph.py index f1103aea3c..de0d77ba5f 100644 --- a/aiida/common/graph.py +++ b/aiida/common/graph.py @@ -9,7 +9,8 @@ ########################################################################### import os, tempfile -def draw_graph(origin_node, ancestor_depth=None, descendant_depth=None, format='dot', + +def draw_graph(origin_node, ancestor_depth=None, descendant_depth=None, format='dot', filename_suffix=None, include_calculation_inputs=False, include_calculation_outputs=False): """ The algorithm starts from the original node and goes both input-ward and output-ward via a breadth-first algorithm. @@ -184,8 +185,12 @@ def draw_link_settings(inp_id, out_id, link_label, link_type): fout.write("}\n") # Now I am producing the output file - output_file_name = "{0}.{format}".format(origin_node.pk, format=format) - exit_status = os.system('dot -T{format} {0} -o {1}'.format(fname, output_file_name, format=format)) + output_file_name = str(origin_node.pk) + if filename_suffix is not None: + output_file_name += "_{}".format(filename_suffix) + + full_output_file_name = "{0}.{format}".format(output_file_name, format=format) + exit_status = os.system('dot -T{format} {0} -o {1}'.format(fname, full_output_file_name, format=format)) # cleaning up by removing the temporary file os.remove(fname) - return exit_status, output_file_name + return exit_status, full_output_file_name diff --git a/aiida/orm/importexport.py b/aiida/orm/importexport.py index 94c9c229b8..1c4103ddad 100644 --- a/aiida/orm/importexport.py +++ b/aiida/orm/importexport.py @@ -1621,6 +1621,7 @@ def serialize_dict(datadict, remove_fields=[], rename_fields={}, 'ctime', 'dbcomputer', 'label', 'type'], } + def fill_in_query(partial_query, originating_entity_str, current_entity_str, tag_suffixes=[], entity_separator="_"): """ @@ -1689,79 +1690,32 @@ def fill_in_query(partial_query, originating_entity_str, current_entity_str, new_tag_suffixes) -def export_tree(what, folder,allowed_licenses=None, forbidden_licenses=None, - silent=False, input_forward=False, create_reversed=True, - return_reversed=False, call_reversed=False, **kwargs): +def node_export_set_expansion(to_be_visited, input_forward=False, create_reversed=True, return_reversed=False, + call_reversed=False): """ - Export the entries passed in the 'what' list to a file tree. - :todo: limit the export to finished or failed calculations. - :param what: a list of entity instances; they can belong to - different models/entities. - :param folder: a :py:class:`Folder ` object - :param input_forward: Follow forward INPUT links (recursively) when - calculating the node set to export. - :param create_reversed: Follow reversed CREATE links (recursively) when - calculating the node set to export. - :param return_reversed: Follow reversed RETURN links (recursively) when - calculating the node set to export. - :param call_reversed: Follow reversed CALL links (recursively) when - calculating the node set to export. - :param allowed_licenses: a list or a function. If a list, then checks - whether all licenses of Data nodes are in the list. If a function, - then calls function for licenses of Data nodes expecting True if - license is allowed, False otherwise. - :param forbidden_licenses: a list or a function. If a list, then checks - whether all licenses of Data nodes are in the list. If a function, - then calls function for licenses of Data nodes expecting True if - license is allowed, False otherwise. - :param silent: suppress debug prints - :raises LicensingException: if any node is licensed under forbidden - license + This function implements the set expansion rules of the export function. Given a set of nodes, the set is + expanded based on the rules that are described at issue #1102 + + Link type Tail & head of the link Link traversal Export direct successor / direct predecessor + INPUT (Data, Calculation) Forward No Reversed Yes + CREATE (Calculation, Data) Forward Yes Reversed Yes + RETURN (Calculation, Data) Forward Yes Reversed No + CALL (Calculation [caller], Calculation [called]) Forward Yes Reversed No + + :param to_be_visited: The initial set of nodes that should be expanded + :param input_forward: Override the input forward rule + :param create_reversed: Override the create reversed rule + :param return_reversed: Override the return reversed rule + :param call_reversed: Override the call reversed rule + :return: The enriched set of nodes """ - import json - import aiida - - from aiida.orm import Node, Calculation, Data, Group, Code + from aiida.orm import Calculation, Data, Code from aiida.common.links import LinkType - from aiida.common.folders import RepositoryFolder from aiida.orm.querybuilder import QueryBuilder - if not silent: - print "STARTING EXPORT..." - - EXPORT_VERSION = '0.3' - all_fields_info, unique_identifiers = get_all_fields_info() - - # The dictionary that contains nodes ids of nodes that should be visited. - # The nodes ids are categorised per node type - to_be_visited = dict() - to_be_visited[Calculation] = set() - to_be_visited[Data] = set() # The set that contains the nodes ids of the nodes that should be exported to_be_exported = set() - given_node_entry_ids = set() - given_group_entry_ids = set() - - # I keep the pks of the given entities - for entry in what: - if issubclass(entry.__class__, Group): - given_group_entry_ids.add(entry.pk) - elif issubclass(entry.__class__, Node): - # The Code node should be treated as a Data node - if (issubclass(entry.__class__, Data) or - issubclass(entry.__class__, Code)): - to_be_visited[Data].add(entry.pk) - elif issubclass(entry.__class__, Calculation): - to_be_visited[Calculation].add(entry.pk) - else: - raise ValueError("I was given {}, which is not a Node nor Group " - "instance. It is of type {}" - .format(entry, entry.__class__)) - - # We will iteratively explore the AiiDA graph to find further nodes that - # should also be exported. - # We repeat until there are no further nodes to be visited while to_be_visited[Calculation] or to_be_visited[Data]: # If is is a calculation node @@ -1795,30 +1749,6 @@ def export_tree(what, folder,allowed_licenses=None, forbidden_licenses=None, res = {_[0] for _ in qb.all()} to_be_visited[Data].update(res - to_be_exported) - - # INPUT(Data, Calculation) - Forward - if input_forward: - qb = QueryBuilder() - qb.append(Data, tag='predecessor', project=['id'], - filters={'id': {'==': curr_node_id}}) - qb.append(Calculation, output_of='predecessor', - edge_filters={ - 'type': { - '==': LinkType.INPUT.value}}) - res = {_[0] for _ in qb.all()} - to_be_visited[Data].update(res - to_be_exported) - # The same until Code becomes a subclass of Data - qb = QueryBuilder() - qb.append(Code, tag='predecessor', project=['id'], - filters={'id': {'==': curr_node_id}}) - qb.append(Calculation, output_of='predecessor', - edge_filters={ - 'type': { - '==': LinkType.INPUT.value}}) - res = {_[0] for _ in qb.all()} - to_be_visited[Data].update(res - to_be_exported) - - # CREATE/RETURN(Calculation, Data) - Forward qb = QueryBuilder() qb.append(Calculation, tag='predecessor', @@ -1842,35 +1772,90 @@ def export_tree(what, folder,allowed_licenses=None, forbidden_licenses=None, res = {_[0] for _ in qb.all()} to_be_visited[Data].update(res - to_be_exported) + # CALL(Calculation, Calculation) - Forward + qb = QueryBuilder() + qb.append(Calculation, tag='predecessor', + filters={'id': {'==': curr_node_id}}) + qb.append(Calculation, output_of='predecessor', project=['id'], + edge_filters={ + 'type': { + '==': LinkType.CALL.value}}) + res = {_[0] for _ in qb.all()} + to_be_visited[Calculation].update(res - to_be_exported) + + # CALL(Calculation, Calculation) - Reversed + if call_reversed: + qb = QueryBuilder() + qb.append(Calculation, tag='predecessor', project=['id']) + qb.append(Calculation, output_of='predecessor', + filters={'id': {'==': curr_node_id}}, + edge_filters={ + 'type': { + '==': LinkType.CALL.value}}) + res = {_[0] for _ in qb.all()} + to_be_visited[Calculation].update(res - to_be_exported) + + # If it is a Data node + else: + curr_node_id = to_be_visited[Data].pop() + # If it is already visited continue to the next node + if curr_node_id in to_be_exported: + continue + # Otherwise say that it is a node to be exported + else: + to_be_exported.add(curr_node_id) + + # INPUT(Data, Calculation) - Forward + if input_forward: + qb = QueryBuilder() + qb.append(Data, tag='predecessor', + filters={'id': {'==': curr_node_id}}) + qb.append(Calculation, output_of='predecessor', + project=['id'], + edge_filters={ + 'type': { + '==': LinkType.INPUT.value}}) + res = {_[0] for _ in qb.all()} + to_be_visited[Calculation].update(res - to_be_exported) + # The same until Code becomes a subclass of Data + qb = QueryBuilder() + qb.append(Code, tag='predecessor', + filters={'id': {'==': curr_node_id}}) + qb.append(Calculation, output_of='predecessor', + project=['id'], + edge_filters={ + 'type': { + '==': LinkType.INPUT.value}}) + res = {_[0] for _ in qb.all()} + to_be_visited[Calculation].update(res - to_be_exported) # CREATE(Calculation, Data) - Reversed if create_reversed: qb = QueryBuilder() - qb.append(Calculation, tag='predecessor') - qb.append(Data, output_of='predecessor', project=['id'], + qb.append(Calculation, tag='predecessor', project=['id']) + qb.append(Data, output_of='predecessor', filters={'id': {'==': curr_node_id}}, edge_filters={ 'type': { - 'in': [LinkType.CREATE.value]}}) + '==': LinkType.CREATE.value}}) res = {_[0] for _ in qb.all()} - to_be_visited[Data].update(res - to_be_exported) + to_be_visited[Calculation].update(res - to_be_exported) # The same until Code becomes a subclass of Data qb = QueryBuilder() - qb.append(Calculation, tag='predecessor') - qb.append(Code, output_of='predecessor', project=['id'], + qb.append(Calculation, tag='predecessor', project=['id']) + qb.append(Code, output_of='predecessor', filters={'id': {'==': curr_node_id}}, edge_filters={ 'type': { - 'in': [LinkType.CREATE.value]}}) + '==': LinkType.CREATE.value}}) res = {_[0] for _ in qb.all()} - to_be_visited[Data].update(res - to_be_exported) - + to_be_visited[Calculation].update(res - to_be_exported) # RETURN(Calculation, Data) - Reversed if return_reversed: qb = QueryBuilder() - qb.append(Calculation, tag='predecessor') - qb.append(Data, output_of='predecessor', project=['id'], + qb.append(Calculation, tag='predecessor', project=['id']) + qb.append(Data, output_of='predecessor', filters={'id': {'==': curr_node_id}}, edge_filters={ 'type': { @@ -1879,116 +1864,398 @@ def export_tree(what, folder,allowed_licenses=None, forbidden_licenses=None, to_be_visited[Data].update(res - to_be_exported) # The same until Code becomes a subclass of Data qb = QueryBuilder() - qb.append(Calculation, tag='predecessor') - qb.append(Code, output_of='predecessor', project=['id'], + qb.append(Calculation, tag='predecessor', project=['id']) + qb.append(Code, output_of='predecessor', filters={'id': {'==': curr_node_id}}, edge_filters={ 'type': { 'in': [LinkType.RETURN.value]}}) res = {_[0] for _ in qb.all()} - to_be_visited[Data].update(res - to_be_exported) + to_be_visited[Calculation].update(res - to_be_exported) + return to_be_exported - # CALL(Calculation, Calculation) - Forward - qb = QueryBuilder() - qb.append(Calculation, tag='predecessor', - filters={'id': {'==': curr_node_id}}) - qb.append(Calculation, output_of='predecessor', project=['id'], - edge_filters={ - 'type': { - '==': LinkType.CALL.value}}) - res = {_[0] for _ in qb.all()} - to_be_visited[Calculation].update(res - to_be_exported) +def link_extraction(all_nodes_pk, input_forward=False, create_reversed=True, return_reversed=False, + call_reversed=False): + """ + This function returns the links that corresponds to the node set expanded by node_export_set_expansion function. + set expansion rules of the export function. + The logic is described at issue #1102 + + Link type, Tail & head of the link, Link export + INPUT, (Data, Calculation), Backward, by the Calculation node, forward too by the Data node, if the user modifies + the default behaviour + CREATE, (Calculation, Data), Forward, by the Calculation node, backward, by the Data node + RETURN, (Calculation, Data), Forward, by the Calculation node + CALL, (Calculation [caller], Calculation [called]), Forward, by the Calculation [caller] node. Also backward too by + the Calculation [called] node, if the user modifies the default behaviour + + :param all_nodes_pk: The nodes that will be exported + :param input_forward: Override the input forward rule + :param create_reversed: Override the create reversed rule + :param return_reversed: Override the return reversed rule + :param call_reversed: Override the call reversed rule + :return: The links that should be exported + """ + from aiida.orm import Calculation, Data, Code + from aiida.common.links import LinkType + from aiida.orm.querybuilder import QueryBuilder - # CALL(Calculation, Calculation) - Reversed - if call_reversed: - qb = QueryBuilder() - qb.append(Calculation, tag='predecessor') - qb.append(Calculation, output_of='predecessor', project=['id'], - filters={'id': {'==': curr_node_id}}, - edge_filters={ - 'type': { - '==': LinkType.CALL.value}}) - res = {_[0] for _ in qb.all()} - to_be_visited[Calculation].update(res - to_be_exported) + links_uuid_dict = dict() + # INPUT (Data, Calculation) - Forward, by the Calculation node + if input_forward: + # INPUT (Data, Calculation) + links_qb = QueryBuilder() + links_qb.append(Data, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Calculation, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.INPUT.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + # INPUT (Code, Calculation) + # The same as above until Code becomes a subclass of Data + links_qb = QueryBuilder() + links_qb.append(Code, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Calculation, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.INPUT.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val - # If it is a Data node - else: - curr_node_id = to_be_visited[Data].pop() - # If it is already visited continue to the next node - if curr_node_id in to_be_exported: - continue - # Otherwise say that it is a node to be exported - else: - to_be_exported.add(curr_node_id) + # INPUT (Data, Calculation) - Backward, by the Calculation node + links_qb = QueryBuilder() + links_qb.append(Data, + project=['uuid'], tag='input') + links_qb.append(Calculation, + project=['uuid'], tag='output', + filters={'id': {'in': all_nodes_pk}}, + edge_filters={'type': {'==': LinkType.INPUT.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + # INPUT (Data, Calculation) - Backward, by the Calculation node + # The same as above until Code becomes a subclass of Data + links_qb = QueryBuilder() + links_qb.append(Code, + project=['uuid'], tag='input') + links_qb.append(Calculation, + project=['uuid'], tag='output', + filters={'id': {'in': all_nodes_pk}}, + edge_filters={ + 'type': {'==': LinkType.INPUT.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + + # CREATE (Calculation, Data) - Forward, by the Calculation node + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Data, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.CREATE.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + + # CREATE (Calculation, Code) - Forward, by the Calculation node + # The same as above until Code becomes a subclass of Data + # This case will not happen (with the current setup - a code is not + # created by a calculation) but it is addded for completeness + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Code, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.CREATE.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val - # Case 2: - # CREATE(Calculation, Data) - Reversed - qb = QueryBuilder() - qb.append(Calculation, tag='predecessor', project=['id']) - qb.append(Data, output_of='predecessor', - filters={'id': {'==': curr_node_id}}, - edge_filters={ - 'type': { - '==': LinkType.CREATE.value}}) - res = {_[0] for _ in qb.all()} - to_be_visited[Calculation].update(res - to_be_exported) - # The same until Code becomes a subclass of Data - qb = QueryBuilder() - qb.append(Calculation, tag='predecessor', project=['id']) - qb.append(Code, output_of='predecessor', - filters={'id': {'==': curr_node_id}}, - edge_filters={ - 'type': { - '==': LinkType.CREATE.value}}) - res = {_[0] for _ in qb.all()} - to_be_visited[Calculation].update(res - to_be_exported) + # CREATE (Calculation, Data) - Backward, by the Data node + if create_reversed: + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Data, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.CREATE.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val - given_node_entry_ids.update(to_be_exported) + # CREATE (Calculation, Code) - Backward, by the Code node + # The same as above until Code becomes a subclass of Data + # This case will not happen (with the current setup - a code is not + # created by a calculation) but it is addded for completeness + if create_reversed: + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Data, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.CREATE.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + + # RETURN (Calculation, Data) - Forward, by the Calculation node + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Data, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.RETURN.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + + # RETURN (Calculation, Data) - Backward, by the Data node + if return_reversed: + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input') + links_qb.append(Data, + project=['uuid'], tag='output', + filters={'id': {'in': all_nodes_pk}}, + edge_filters={'type': {'==': LinkType.RETURN.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + + # CALL (Calculation [caller], Calculation [called]) - Forward, by + # the Calculation node + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input', + filters={'id': {'in': all_nodes_pk}}) + links_qb.append(Calculation, + project=['uuid'], tag='output', + edge_filters={'type': {'==': LinkType.CALL.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + + # CALL (Calculation [caller], Calculation [called]) - Backward, + # by the Calculation [called] node + if call_reversed: + links_qb = QueryBuilder() + links_qb.append(Calculation, + project=['uuid'], tag='input') + links_qb.append(Calculation, + project=['uuid'], tag='output', + filters={'id': {'in': all_nodes_pk}}, + edge_filters={'type': {'==': LinkType.CALL.value}}, + edge_project=['label', 'type'], output_of='input') + for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): + val = { + 'input': str(input_uuid), + 'output': str(output_uuid), + 'label': str(link_label), + 'type': str(link_type) + } + links_uuid_dict[frozenset(val.items())] = val + + return links_uuid_dict.values() + + +def export_tree(what, folder, allowed_licenses=None, forbidden_licenses=None, + silent=False, input_forward=False, create_reversed=True, + return_reversed=False, call_reversed=False, **kwargs): + """ + Export the entries passed in the 'what' list to a file tree. + :todo: limit the export to finished or failed calculations. + :param what: a list of entity instances; they can belong to + different models/entities. + :param folder: a :py:class:`Folder ` object + :param input_forward: Follow forward INPUT links (recursively) when + calculating the node set to export. + :param create_reversed: Follow reversed CREATE links (recursively) when + calculating the node set to export. + :param return_reversed: Follow reversed RETURN links (recursively) when + calculating the node set to export. + :param call_reversed: Follow reversed CALL links (recursively) when + calculating the node set to export. + :param allowed_licenses: a list or a function. If a list, then checks + whether all licenses of Data nodes are in the list. If a function, + then calls function for licenses of Data nodes expecting True if + license is allowed, False otherwise. + :param forbidden_licenses: a list or a function. If a list, then checks + whether all licenses of Data nodes are in the list. If a function, + then calls function for licenses of Data nodes expecting True if + license is allowed, False otherwise. + :param silent: suppress debug prints + :raises LicensingException: if any node is licensed under forbidden + license + """ + import json + import aiida + + from aiida.orm import Node, Calculation, Data, Group, Code + from aiida.common.folders import RepositoryFolder + from aiida.orm.querybuilder import QueryBuilder + if not silent: + print "STARTING EXPORT..." + print("Current values of the export node set expansion rules") + print("input_forward: ", input_forward) + print("create_reversed: ", create_reversed) + print("return_reversed: ", return_reversed) + print("call_reversed: ", call_reversed) + + EXPORT_VERSION = '0.3' + + all_fields_info, unique_identifiers = get_all_fields_info() + + to_be_exported_group_entry_ids = set() + + # The dictionary that contains node ids/pks of given nodes. They will be enriched with the ids/pks + # of related nodes given a set of rules and arguments. + to_be_visited = dict() + to_be_visited[Calculation] = set() + to_be_visited[Data] = set() + + # Getting the ids/pks of the given entities and dividing them + # to the right categories + for entry in what: + if issubclass(entry.__class__, Group): + to_be_exported_group_entry_ids.add(entry.pk) + elif issubclass(entry.__class__, Node): + # The Code node should be treated as a Data node + if (issubclass(entry.__class__, Data) or + issubclass(entry.__class__, Code)): + to_be_visited[Data].add(entry.pk) + elif issubclass(entry.__class__, Calculation): + to_be_visited[Calculation].add(entry.pk) + else: + raise ValueError("I was given {}, which is not a Node nor Group " + "instance. It is of type {}" + .format(entry, entry.__class__)) + + # Search fot the node entities that should be exported based on the set of export rules + # and arguments. + to_be_exported_node_entry_ids = node_export_set_expansion(to_be_visited, input_forward, create_reversed, + return_reversed, call_reversed) # Here we get all the columns that we plan to project per entity that we # would like to extract - given_entities = list() - if len(given_group_entry_ids) > 0: - given_entities.append(GROUP_ENTITY_NAME) - if len(given_node_entry_ids) > 0: - given_entities.append(NODE_ENTITY_NAME) + to_be_exported_entities = list() + if len(to_be_exported_group_entry_ids) > 0: + to_be_exported_entities.append(GROUP_ENTITY_NAME) + if len(to_be_exported_node_entry_ids) > 0: + to_be_exported_entities.append(NODE_ENTITY_NAME) entries_to_add = dict() - for given_entity in given_entities: + for to_be_exported_entity in to_be_exported_entities: project_cols = ["id"] # The following gets a list of fields that we need, # e.g. user, mtime, uuid, computer - entity_prop = all_fields_info[given_entity].keys() + entity_prop = all_fields_info[to_be_exported_entity].keys() # Here we do the necessary renaming of properties for prop in entity_prop: # nprop contains the list of projections - nprop = (file_fields_to_model_fields[given_entity][prop] - if file_fields_to_model_fields[given_entity].has_key(prop) + nprop = (file_fields_to_model_fields[to_be_exported_entity][prop] + if file_fields_to_model_fields[to_be_exported_entity].has_key(prop) else prop) project_cols.append(nprop) # Getting the ids that correspond to the right entity - entry_ids_to_add = (given_node_entry_ids - if (given_entity == NODE_ENTITY_NAME) - else given_group_entry_ids) + entry_ids_to_add = (to_be_exported_node_entry_ids + if (to_be_exported_entity == NODE_ENTITY_NAME) + else to_be_exported_group_entry_ids) qb = QueryBuilder() - qb.append(entity_names_to_entities[given_entity], + qb.append(entity_names_to_entities[to_be_exported_entity], filters={"id": {"in": entry_ids_to_add}}, project=project_cols, - tag=given_entity, outerjoin=True) - entries_to_add[given_entity] = qb + tag=to_be_exported_entity, outerjoin=True) + entries_to_add[to_be_exported_entity] = qb # TODO (Spyros) To see better! Especially for functional licenses # Check the licenses of exported data. if allowed_licenses is not None or forbidden_licenses is not None: qb = QueryBuilder() qb.append(Node, project=["id", "attributes.source.license"], - filters={"id": {"in": given_node_entry_ids}}) + filters={"id": {"in": to_be_exported_node_entry_ids}}) # Skip those nodes where the license is not set (this is the standard behavior with Django) node_licenses = list((a,b) for [a,b] in qb.all() if b is not None) check_licences(node_licenses, allowed_licenses, forbidden_licenses) @@ -2073,239 +2340,12 @@ def export_tree(what, folder,allowed_licenses=None, forbidden_licenses=None, if not silent: print "STORING NODE LINKS..." - links_uuid_dict = dict() + # Get alla the links that correspond to the selected nodes if len(all_nodes_pk) > 0: - # INPUT (Data, Calculation) - Forward, by the Calculation node - if input_forward: - # INPUT (Data, Calculation) - links_qb = QueryBuilder() - links_qb.append(Data, - project=['uuid'], tag='input', - filters = {'id': {'in': all_nodes_pk}}) - links_qb.append(Calculation, - project=['uuid'], tag='output', - edge_filters={'type':{'==':LinkType.INPUT.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - # INPUT (Code, Calculation) - # The same as above until Code becomes a subclass of Data - links_qb = QueryBuilder() - links_qb.append(Code, - project=['uuid'], tag='input', - filters = {'id': {'in': all_nodes_pk}}) - links_qb.append(Calculation, - project=['uuid'], tag='output', - edge_filters={'type':{'==':LinkType.INPUT.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - # INPUT (Data, Calculation) - Backward, by the Calculation node - links_qb = QueryBuilder() - links_qb.append(Data, - project=['uuid'], tag='input') - links_qb.append(Calculation, - project=['uuid'], tag='output', - filters={'id': {'in': all_nodes_pk}}, - edge_filters={'type':{'==':LinkType.INPUT.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - # INPUT (Data, Calculation) - Backward, by the Calculation node - # The same as above until Code becomes a subclass of Data - links_qb = QueryBuilder() - links_qb.append(Code, - project=['uuid'], tag='input') - links_qb.append(Calculation, - project=['uuid'], tag='output', - filters={'id': {'in': all_nodes_pk}}, - edge_filters={ - 'type': {'==': LinkType.INPUT.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type': str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - # CREATE (Calculation, Data) - Forward, by the Calculation node - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input', - filters={'id': {'in': all_nodes_pk}}) - links_qb.append(Data, - project=['uuid'], tag='output', - edge_filters={'type': {'==': LinkType.CREATE.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - # CREATE (Calculation, Code) - Forward, by the Calculation node - # The same as above until Code becomes a subclass of Data - # This case will not happen (with the current setup - a code is not - # created by a calculation) but it is addded for completeness - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input', - filters={'id': {'in': all_nodes_pk}}) - links_qb.append(Code, - project=['uuid'], tag='output', - edge_filters={'type': {'==': LinkType.CREATE.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - - # CREATE (Calculation, Data) - Backward, by the Data node - if create_reversed: - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input', - filters={'id': {'in': all_nodes_pk}}) - links_qb.append(Data, - project=['uuid'], tag='output', - edge_filters={'type': {'==': LinkType.CREATE.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - # CREATE (Calculation, Code) - Backward, by the Code node - # The same as above until Code becomes a subclass of Data - # This case will not happen (with the current setup - a code is not - # created by a calculation) but it is addded for completeness - if create_reversed: - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input', - filters={'id': {'in': all_nodes_pk}}) - links_qb.append(Data, - project=['uuid'], tag='output', - edge_filters={'type': {'==': LinkType.CREATE.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - # RETURN (Calculation, Data) - Forward, by the Calculation node - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input', - filters={'id': {'in': all_nodes_pk}}) - links_qb.append(Data, - project=['uuid'], tag='output', - edge_filters={'type': {'==': LinkType.RETURN.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - # RETURN (Calculation, Data) - Backward, by the Data node - if return_reversed: - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input') - links_qb.append(Data, - project=['uuid'], tag='output', - filters={'id': {'in': all_nodes_pk}}, - edge_filters={'type': {'==': LinkType.RETURN.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - # CALL (Calculation [caller], Calculation [called]) - Forward, by - # the Calculation node - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input', - filters={'id': {'in': all_nodes_pk}}) - links_qb.append(Calculation, - project=['uuid'], tag='output', - edge_filters={'type': {'==': LinkType.CALL.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - # CALL (Calculation [caller], Calculation [called]) - Backward, - # by the Calculation [called] node - if call_reversed: - links_qb = QueryBuilder() - links_qb.append(Calculation, - project=['uuid'], tag='input') - links_qb.append(Calculation, - project=['uuid'], tag='output', - filters={'id': {'in': all_nodes_pk}}, - edge_filters={'type': {'==': LinkType.CALL.value}}, - edge_project=['label', 'type'], output_of='input') - for input_uuid, output_uuid, link_label, link_type in links_qb.iterall(): - val = { - 'input': str(input_uuid), - 'output': str(output_uuid), - 'label': str(link_label), - 'type':str(link_type) - } - links_uuid_dict[frozenset(val.items())] = val - - links_uuid = links_uuid_dict.values() + links_uuid = link_extraction(all_nodes_pk, input_forward=False, create_reversed=True, return_reversed=False, + call_reversed=False) + else: + links_uuid = list() if not silent: print "STORING GROUP ELEMENTS..." diff --git a/docs/source/nitpick-exceptions b/docs/source/nitpick-exceptions index 9704fef531..812314429b 100644 --- a/docs/source/nitpick-exceptions +++ b/docs/source/nitpick-exceptions @@ -152,3 +152,4 @@ py:class aiida.work.interstep.Builder py:class Builder py:class Data py:meth aiida.orm.group.Group.get_from_string +py:class _WorkChainSpec \ No newline at end of file