Export/import of extras:

aiida/backends/tests/export_and_import.py

@@ -2458,3 +2458,154 @@ def test_multiple_user_comments_for_single_node(self):
 
         finally:
             shutil.rmtree(tmp_folder, ignore_errors=True)
+
+class TestExtras(AiidaTestCase):
+    """Test extras import"""
+
+    @classmethod
+    def setUpClass(cls, *args, **kwargs):
+        """Only run to prepare an export file"""
+        import os 
+        import tempfile
+        from aiida.orm import Data
+        from aiida.orm.importexport import export
+        super(TestExtras, cls).setUpClass()
+        d = Data()
+        d.label = 'my_test_data_node'
+        d.store()
+        d.set_extras({'b': 2, 'c': 3})
+        tmp_folder = tempfile.mkdtemp()
+        cls.export_file = os.path.join(tmp_folder, 'export.aiida')
+        export([d], outfile=cls.export_file, silent=True)
+
+    def setUp(self):
+        """This function runs before every test execution"""
+        from aiida.orm.importexport import import_data
+        self.clean_db()
+        self.insert_data()
+
+    def import_extras(self, mode_new='import'):
+        """Import an aiida database"""
+        from aiida.orm import Data
+        from aiida.orm.querybuilder import QueryBuilder
+        import_data(self.export_file, silent=True, extras_mode_new=mode_new)
+        q = QueryBuilder().append(Data, filters={'label': 'my_test_data_node'})
+        self.assertEquals(q.count(), 1)
+        self.imported_node = q.all()[0][0]
+
+    def modify_extras(self, mode_existing):
+        """Import the same aiida database again"""
+        from aiida.orm import Data
+        from aiida.orm.querybuilder import QueryBuilder
+        self.imported_node.set_extra('a', 1)
+        self.imported_node.set_extra('b', 1000)
+        self.imported_node.del_extra('c')
+
+        import_data(self.export_file, silent=True, extras_mode_existing=mode_existing)
+
+        # Query again the database
+        q = QueryBuilder().append(Data, filters={'label': 'my_test_data_node'})
+        self.assertEquals(q.count(), 1)
+        return q.all()[0][0]
+
+    def tearDown(self):
+        pass
+
+
+    def test_import_of_extras(self):
+        """Check if extras are properly imported"""
+        self.import_extras()
+        self.assertEquals(self.imported_node.get_extra('b'), 2)
+        self.assertEquals(self.imported_node.get_extra('c'), 3)
+
+    def test_absence_of_extras(self):
+        """Check whether extras are not imported if the mode is set to 'none'"""
+        self.import_extras(mode_new='none')
+        with self.assertRaises(AttributeError):
+            # the extra 'b' should not exist
+            self.imported_node.get_extra('b')
+        with self.assertRaises(AttributeError):
+            # the extra 'c' should not exist
+            self.imported_node.get_extra('c')
+
+    def test_extras_import_mode_keep_existing(self):
+        """Check if old extras are not modified in case of name collision"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='kcl')
+
+        # Check that extras are imported correctly
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        self.assertEquals(imported_node.get_extra('b'), 1000)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+
+    def test_extras_import_mode_update_existing(self):
+        """Check if old extras are modified in case of name collision"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='kcu')
+
+        # Check that extras are imported correctly
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        self.assertEquals(imported_node.get_extra('b'), 2)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+
+    def test_extras_import_mode_mirror(self):
+        """Check if old extras are fully overwritten by the imported ones"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='ncu')
+
+        # Check that extras are imported correctly
+        with self.assertRaises(AttributeError): # the extra 
+            # 'a' should not exist, as the extras were fully mirrored with respect to
+            # the imported node
+            imported_node.get_extra('a')
+        self.assertEquals(imported_node.get_extra('b'), 2)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+
+    def test_extras_import_mode_none(self):
+        """Check if old extras are fully overwritten by the imported ones"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='knl')
+
+        # Check if extras are imported correctly
+        self.assertEquals(imported_node.get_extra('b'), 1000)
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        with self.assertRaises(AttributeError): # the extra
+            # 'c' should not exist, as the extras were keept untached
+            imported_node.get_extra('c')
+
+    def test_extras_import_mode_strange(self):
+        """Check a mode that is probably does not make much sense but is still available"""
+        self.import_extras()
+        imported_node = self.modify_extras(mode_existing='kcd')
+
+        # Check if extras are imported correctly
+        self.assertEquals(imported_node.get_extra('a'), 1)
+        self.assertEquals(imported_node.get_extra('c'), 3)
+        with self.assertRaises(AttributeError): # the extra
+            # 'b' should not exist, as the collided extras are deleted
+            imported_node.get_extra('b')
+
+    def test_extras_import_mode_correct(self):
+        """Test all possible import modes except 'ask' """
+        self.import_extras()
+        for l1 in ['k', 'n']: # keep or not keep old extras
+            for l2 in ['n', 'c']: # create or not create new extras
+                for l3 in ['l', 'u', 'd']: # leave old, update or delete collided extras
+                    mode = l1 + l2 + l3
+                    import_data(self.export_file, silent=True, extras_mode_existing=mode)
+
+    def test_extras_import_mode_wrong(self):
+        """Check a mode that is wrong"""
+        self.import_extras()
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='xnd') # first letter is wrong
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='nxd') # second letter is wrong
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='nnx') # third letter is wrong
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='n') # too short
+        with self.assertRaises(ValueError):
+            import_data(self.export_file, silent=True, extras_mode_existing='nndnn') # too long
+        with self.assertRaises(TypeError):
+            import_data(self.export_file, silent=True, extras_mode_existing=5) # wrong type

aiida/cmdline/commands/cmd_import.py

@@ -11,6 +11,7 @@
 from __future__ import division
 from __future__ import print_function
 from __future__ import absolute_import
+from enum import Enum
 import click
 
 from aiida.cmdline.commands.cmd_verdi import verdi
@@ -19,6 +20,19 @@
 from aiida.cmdline.utils import decorators, echo
 from aiida.common import exceptions
 
+EXTRAS_MODE_EXISTING = ['keep_existing', 'update_existing', 'mirror', 'none', 'ask']
+EXTRAS_MODE_NEW = ['import', 'none']
+
+
+# pylint: disable=too-few-public-methods
+class ExtrasImportCode(Enum):
+    """Exit codes for the verdi command line."""
+    keep_existing = 'kcl'
+    update_existing = 'kcu'
+    mirror = 'ncu'
+    none = 'knl'
+    ask = 'kca'
+
 
 @verdi.command('import')
 @click.argument('archives', nargs=-1, type=click.Path(exists=True, readable=True))
@@ -35,8 +49,28 @@
     type=GroupParamType(create_if_not_exist=True),
     help='Specify group to which all the import nodes will be added. If such a group does not exist, it will be'
     ' created automatically.')
+@click.option(
+    '-e',
+    '--extras-mode-existing',
+    type=click.Choice(EXTRAS_MODE_EXISTING),
+    default='keep_existing',
+    help="Specify which extras from the export archive should be imported for nodes that are already contained in the "
+    "database: "
+    "ask: import all extras and prompt what to do for existing extras. "
+    "keep_existing: import all extras and keep original value of existing extras. "
+    "update_existing: import all extras and overwrite value of existing extras. "
+    "mirror: import all extras and remove any existing extras that are not present in the archive. "
+    "none: do not import any extras.")
+@click.option(
+    '-n',
+    '--extras-mode-new',
+    type=click.Choice(EXTRAS_MODE_NEW),
+    default='import',
+    help="Specify whether to import extras of new nodes: "
+    "import: import extras. "
+    "none: do not import extras.")
 @decorators.with_dbenv()
-def cmd_import(archives, webpages, group):
+def cmd_import(archives, webpages, group, extras_mode_existing, extras_mode_new):
     """Import one or multiple exported AiiDA archives
 
     The ARCHIVES can be specified by their relative or absolute file path, or their HTTP URL.
@@ -78,7 +112,11 @@ def cmd_import(archives, webpages, group):
         echo.echo_info('importing archive {}'.format(archive))
 
         try:
-            import_data(archive, group)
+            import_data(
+                archive,
+                group,
+                extras_mode_existing=ExtrasImportCode[extras_mode_existing].value,
+                extras_mode_new=extras_mode_new)
         except exceptions.IncompatibleArchiveVersionError as exception:
             echo.echo_warning('{} cannot be imported: {}'.format(archive, exception))
             echo.echo_warning('run `verdi export migrate {}` to update it'.format(archive))
@@ -105,7 +143,11 @@ def cmd_import(archives, webpages, group):
             echo.echo_success('archive downloaded, proceeding with import')
 
             try:
-                import_data(temp_folder.get_abs_path(temp_file), group)
+                import_data(
+                    temp_folder.get_abs_path(temp_file),
+                    group,
+                    extras_mode_existing=ExtrasImportCode[extras_mode_existing].value,
+                    extras_mode_new=extras_mode_new)
             except exceptions.IncompatibleArchiveVersionError as exception:
                 echo.echo_warning('{} cannot be imported: {}'.format(archive, exception))
                 echo.echo_warning('download the archive file and run `verdi export migrate` to update it')