add --in-place flag to verdi export migrate

aiida/cmdline/commands/cmd_export.py

@@ -11,6 +11,7 @@
 """`verdi export` command."""
 
 import os
+import tempfile
 
 import click
 import tabulate
@@ -141,31 +142,45 @@ def create(
 
 @verdi_export.command('migrate')
 @arguments.INPUT_FILE()
-@arguments.OUTPUT_FILE()
+@arguments.OUTPUT_FILE(required=False)
 @options.ARCHIVE_FORMAT()
 @options.FORCE(help='overwrite output file if it already exists')
+@click.option('-i', '--in-place', is_flag=True, help='Migrate the archive in place, overwriting the original file.')
 @options.SILENT()
 @click.option(
     '-v',
     '--version',
     type=click.STRING,
     required=False,
     metavar='VERSION',
-    help='Specify an exact archive version to migrate to. By default the most recent version is taken.'
+    # Note: Adding aiida.tools.EXPORT_VERSION as a default value explicitly would result in a slow import of
+    # aiida.tools and, as a consequence, aiida.orm. As long as this is the case, better determine the latest export
+    # version inside the function when needed.
+    help='Archive format version to migrate to (defaults to latest version).',
 )
-def migrate(input_file, output_file, force, silent, archive_format, version):
+def migrate(input_file, output_file, force, silent, in_place, archive_format, version):
     # pylint: disable=too-many-locals,too-many-statements,too-many-branches
     """Migrate an export archive to a more recent format version."""
     import tarfile
     import zipfile
 
     from aiida.common import json
     from aiida.common.folders import SandboxFolder
-    from aiida.tools.importexport import EXPORT_VERSION, migration, extract_zip, extract_tar, ArchiveMigrationError
+    from aiida.tools.importexport import migration, extract_zip, extract_tar, ArchiveMigrationError, EXPORT_VERSION
 
     if version is None:
         version = EXPORT_VERSION
 
+    if in_place:
+        if output_file:
+            echo.echo_critical('output file specified together with --in-place flag')
+        tempdir = tempfile.TemporaryDirectory()
+        output_file = os.path.join(tempdir.name, 'archive.aiida')
+    elif not output_file:
+        echo.echo_critical(
+            'no output file specified. Please add --in-place flag if you would like to migrate in place.'
+        )
+
     if os.path.exists(output_file) and not force:
         echo.echo_critical('the output file already exists')
 
@@ -187,6 +202,10 @@ def migrate(input_file, output_file, force, silent, archive_format, version):
             echo.echo_critical('export archive does not contain the required file {}'.format(fhandle.filename))
 
         old_version = migration.verify_metadata_version(metadata)
+        if version <= old_version:
+            echo.echo_success('nothing to be done - archive already at version {} >= {}'.format(old_version, version))
+            return
+
         try:
             new_version = migration.migrate_recursively(metadata, data, folder, version)
         except ArchiveMigrationError as exception:
@@ -212,5 +231,9 @@ def migrate(input_file, output_file, force, silent, archive_format, version):
             with tarfile.open(output_file, 'w:gz', format=tarfile.PAX_FORMAT, dereference=True) as archive:
                 archive.add(folder.abspath, arcname='')
 
+        if in_place:
+            os.rename(output_file, input_file)
+            tempdir.cleanup()
+
         if not silent:
             echo.echo_success('migrated the archive from version {} to {}'.format(old_version, new_version))

aiida/tools/importexport/migration/__init__.py

@@ -52,8 +52,8 @@ def migrate_recursively(metadata, data, folder, version=EXPORT_VERSION):
 
     try:
         if old_version == version:
-            raise ArchiveMigrationError('Your export file is already at the version {}'.format(version))
-        elif old_version > version:
+            return old_version
+        if old_version > version:
             raise ArchiveMigrationError('Backward migrations are not supported')
         elif old_version in MIGRATE_FUNCTIONS:
             MIGRATE_FUNCTIONS[old_version](metadata, data, folder)

tests/cmdline/commands/test_export.py

@@ -181,18 +181,6 @@ def test_migrate_version_specific(self):
         finally:
             delete_temporary_file(filename_output)
 
-    def test_migrate_versions_recent(self):
-        """Migrating an archive with the current version should exit with non-zero status."""
-        filename_input = get_archive_file(self.newest_archive, filepath=self.fixture_archive)
-        filename_output = next(tempfile._get_candidate_names())  # pylint: disable=protected-access
-
-        try:
-            options = [filename_input, filename_output]
-            result = self.cli_runner.invoke(cmd_export.migrate, options)
-            self.assertIsNotNone(result.exception)
-        finally:
-            delete_temporary_file(filename_output)
-
     def test_migrate_force(self):
         """Test that passing the -f/--force option will overwrite the output file even if it exists."""
         filename_input = get_archive_file(self.penultimate_archive, filepath=self.fixture_archive)
@@ -213,6 +201,39 @@ def test_migrate_force(self):
                 self.assertTrue(os.path.isfile(filename_output))
                 self.assertEqual(zipfile.ZipFile(filename_output).testzip(), None)
 
+    def test_migrate_in_place(self):
+        """Test that passing the -i/--in-place option will overwrite the passed file."""
+        archive = 'export_v0.1_simple.aiida'
+        target_version = '0.2'
+        filename_input = get_archive_file(archive, filepath=self.fixture_archive)
+        filename_tmp = next(tempfile._get_candidate_names())  # pylint: disable=protected-access
+
+        try:
+            # copy file (don't want to overwrite test data)
+            shutil.copy(filename_input, filename_tmp)
+
+            # specifying both output and in-place should except
+            options = [filename_tmp, '--in-place', '--output-file', 'test.aiida']
+            result = self.cli_runner.invoke(cmd_export.migrate, options)
+            self.assertIsNotNone(result.exception, result.output)
+
+            # specifying neither output nor in-place should except
+            options = [filename_tmp]
+            result = self.cli_runner.invoke(cmd_export.migrate, options)
+            self.assertIsNotNone(result.exception, result.output)
+
+            # check that in-place migration produces a valid archive in place of the old file
+            options = [filename_tmp, '--in-place', '--version', target_version]
+            result = self.cli_runner.invoke(cmd_export.migrate, options)
+            self.assertIsNone(result.exception, result.output)
+            self.assertTrue(os.path.isfile(filename_tmp))
+            # check that files in zip file are ok
+            self.assertEqual(zipfile.ZipFile(filename_tmp).testzip(), None)
+            with Archive(filename_tmp) as archive_object:
+                self.assertEqual(archive_object.version_format, target_version)
+        finally:
+            os.remove(filename_tmp)
+
     def test_migrate_silent(self):
         """Test that the captured output is an empty string when the -s/--silent option is passed."""
         filename_input = get_archive_file(self.penultimate_archive, filepath=self.fixture_archive)

tests/tools/importexport/migration/test_migration.py

@@ -114,9 +114,7 @@ def test_migrate_recursively_specific_version(self):
             with self.assertRaises(ArchiveMigrationError):
                 migrate_recursively(archive.meta_data, archive.data, None, version='0.2')
 
-            # Same version will also raise
-            with self.assertRaises(ArchiveMigrationError):
-                migrate_recursively(archive.meta_data, archive.data, None, version='0.3')
+            migrate_recursively(archive.meta_data, archive.data, None, version='0.3')
 
             migrated_version = '0.5'
             version = migrate_recursively(archive.meta_data, archive.data, None, version=migrated_version)
@@ -186,17 +184,11 @@ def test_wrong_versions(self):
                 )
 
     def test_migrate_newest_version(self):
-        """Test that an exception is raised when an export file with the newest export version is migrated."""
+        """Test that  migrating the latest version runs without complaints."""
         metadata = {'export_version': newest_version}
 
-        with self.assertRaises(ArchiveMigrationError):
-            new_version = migrate_recursively(metadata, {}, None)
-
-            self.assertIsNone(
-                new_version,
-                msg='migrate_recursively should not return anything, '
-                "hence the 'return' should be None, but instead it is {}".format(new_version)
-            )
+        new_version = migrate_recursively(metadata, {}, None)
+        self.assertEqual(new_version, newest_version)
 
     @with_temp_dir
     def test_v02_to_newest(self, temp_dir):