Update SqlAlchemy import and export to include link types

aiidateam · Oct 13, 2017 · 788d320 · 788d320
1 parent 2dbc4a9
commit 788d320
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/aiida/orm/importexport.py b/aiida/orm/importexport.py
@@ -1034,6 +1034,7 @@ def import_data_sqla(in_path, ignore_unknown_nodes=False, silent=False):
     from aiida.utils import timezone
 
     from aiida.orm import Node, Group
+    from aiida.common.links import LinkType
     from aiida.common.folders import SandboxFolder, RepositoryFolder
     from aiida.common.utils import get_object_from_string
     from aiida.common.datastructures import calc_states
@@ -1043,7 +1044,7 @@ def import_data_sqla(in_path, ignore_unknown_nodes=False, silent=False):
     from aiida.backends.sqlalchemy.models.node import DbCalcState
 
     # This is the export version expected by this function
-    expected_export_version = '0.2'
+    expected_export_version = '0.3'
 
     # The name of the subfolder in which the node files are stored
     nodes_export_subfolder = 'nodes'
@@ -1508,7 +1509,7 @@ def import_data_sqla(in_path, ignore_unknown_nodes=False, silent=False):
                         # New link
                         links_to_store.append(DbLink(
                             input_id=in_id, output_id=out_id,
-                            label=link['label']))
+                            label=link['label'], type=LinkType(link['type']).value))
                         if entity_names_to_signatures[
                             LINK_ENTITY_NAME] not in ret_dict:
                             ret_dict[entity_names_to_signatures[
@@ -2064,7 +2065,7 @@ def export_tree_sqla(what, folder, also_parents=True, also_calc_outputs=True,
     if not silent:
         print "STARTING EXPORT..."
 
-    EXPORT_VERSION = '0.2'
+    EXPORT_VERSION = '0.3'
 
     all_fields_info, unique_identifiers = get_all_fields_info_sqla()
 
@@ -2239,13 +2240,14 @@ def export_tree_sqla(what, folder, also_parents=True, also_calc_outputs=True,
         links_qb.append(Node,
                         project=['uuid'], tag='output',
                         filters={'id': {'in': all_nodes_pk}},
-                        edge_project=['label'], output_of='input')
+                        edge_project=['label', 'type'], output_of='input')
 
-        for input_uuid, output_uuid, link_label in links_qb.iterall():
+        for input_uuid, output_uuid, link_label, link_type in links_qb.iterall():
             links_uuid.append({
                 'input': str(input_uuid),
                 'output': str(output_uuid),
-                'label': str(link_label)
+                'label': str(link_label),
+                'type':str(link_type)
             })
 
     if not silent: