Add a progress bar for export and import related functionality (#3599)

When provenance graphs become large, exporting and importing the
resulting archives can become time consuming operations. To give the
user more feedback of how the operation is progressing, a progress bar
is added. The implementation is provided by the `tqdm` library which is
therefore added as a new dependency.

Since the export and import code is still not written in a modular way,
to allow the progress bar to get access to the inner parts and provide
information that is granular enough, it is implemented as a global
singleton that the export and import functions fetch. This is to prevent
having to pass the progress bar instance around in method calls. This is
still not an ideal solution and in the future, this should be replaced
with hooks that methods can call in order to update their status.

A similar problem is faced with the logging of textual progress in the
export and import functions that are done haphazardly with print
statements, even though it concerns module functions and not CLI facing
code. This has been changed to go through logging instead. The log level
is temporarily updated based on the `silent` argument that the export
and import functions take.

Finally, the function signature of `export` and `export_tree` have been
changed. The `what` and `outfile` arguments have been deprecated and
replaced by `entities` and `filename`, respectively. A deprecation
message is printed if they are used and they should be removed with the
release of `aiida-core==2.0.0`.

Author: CasperWA

aiida/cmdline/commands/cmd_export.py

@@ -93,7 +93,8 @@ def create(
     their provenance, according to the rules outlined in the documentation.
     You can modify some of those rules using options of this command.
     """
-    from aiida.tools.importexport import export, export_zip
+    from aiida.tools.importexport import export, ExportFileFormat
+    from aiida.tools.importexport.common.exceptions import ArchiveExportError
 
     entities = []
 
@@ -122,19 +123,18 @@ def create(
     }
 
     if archive_format == 'zip':
-        export_function = export_zip
+        export_format = ExportFileFormat.ZIP
         kwargs.update({'use_compression': True})
     elif archive_format == 'zip-uncompressed':
-        export_function = export_zip
+        export_format = ExportFileFormat.ZIP
         kwargs.update({'use_compression': False})
     elif archive_format == 'tar.gz':
-        export_function = export
+        export_format = ExportFileFormat.TAR_GZIPPED
 
     try:
-        export_function(entities, outfile=output_file, **kwargs)
-
-    except IOError as exception:
-        echo.echo_critical('failed to write the export archive file: {}'.format(exception))
+        export(entities, filename=output_file, file_format=export_format, **kwargs)
+    except ArchiveExportError as exception:
+        echo.echo_critical('failed to write the archive file. Exception: {}'.format(exception))
     else:
         echo.echo_success('wrote the export archive file to {}'.format(output_file))
 

aiida/cmdline/commands/cmd_import.py

@@ -12,6 +12,7 @@
 from enum import Enum
 import traceback
 import urllib.request
+
 import click
 
 from aiida.cmdline.commands.cmd_verdi import verdi
@@ -34,6 +35,45 @@ class ExtrasImportCode(Enum):
     ask = 'kca'
 
 
+def _echo_error(  # pylint: disable=unused-argument
+    message, non_interactive, more_archives, raised_exception, **kwargs
+):
+    """Utility function to help write an error message for ``verdi import``
+
+    :param message: Message following red-colored, bold "Error:".
+    :type message: str
+    :param non_interactive: Whether or not the user should be asked for input for any reason.
+    :type non_interactive: bool
+    :param more_archives: Whether or not there are more archives to import.
+    :type more_archives: bool
+    :param raised_exception: Exception raised during error.
+    :type raised_exception: `Exception`
+    """
+    from aiida.tools.importexport import close_progress_bar, IMPORT_LOGGER
+
+    # Close progress bar, if it exists
+    close_progress_bar(leave=False)
+
+    IMPORT_LOGGER.debug('%s', traceback.format_exc())
+
+    exception = '{}: {}'.format(raised_exception.__class__.__name__, str(raised_exception))
+
+    echo.echo_error(message)
+    echo.echo(exception)
+
+    if more_archives:
+        # There are more archives to go through
+        if non_interactive:
+            # Continue to next archive
+            pass
+        else:
+            # Ask if one should continue to next archive
+            click.confirm('Do you want to continue?', abort=True)
+    else:
+        # There are no more archives
+        click.Abort()
+
+
 def _try_import(migration_performed, file_to_import, archive, group, migration, non_interactive, **kwargs):
     """Utility function for `verdi import` to try to import archive
 
@@ -66,8 +106,12 @@ def _try_import(migration_performed, file_to_import, archive, group, migration,
     except IncompatibleArchiveVersionError as exception:
         if migration_performed:
             # Migration has been performed, something is still wrong
-            crit_message = '{} has been migrated, but it still cannot be imported.\n{}'.format(archive, exception)
-            echo.echo_critical(crit_message)
+            _echo_error(
+                '{} has been migrated, but it still cannot be imported'.format(archive),
+                non_interactive=non_interactive,
+                raised_exception=exception,
+                **kwargs
+            )
         else:
             # Migration has not yet been tried.
             if migration:
@@ -85,18 +129,20 @@ def _try_import(migration_performed, file_to_import, archive, group, migration,
             else:
                 # Abort
                 echo.echo_critical(str(exception))
-    except Exception:
-        echo.echo_error('an exception occurred while importing the archive {}'.format(archive))
-        echo.echo(traceback.format_exc())
-        if not non_interactive:
-            click.confirm('do you want to continue?', abort=True)
+    except Exception as exception:
+        _echo_error(
+            'an exception occurred while importing the archive {}'.format(archive),
+            non_interactive=non_interactive,
+            raised_exception=exception,
+            **kwargs
+        )
     else:
         echo.echo_success('imported archive {}'.format(archive))
 
     return migrate_archive
 
 
-def _migrate_archive(ctx, temp_folder, file_to_import, archive, non_interactive, **kwargs):  # pylint: disable=unused-argument
+def _migrate_archive(ctx, temp_folder, file_to_import, archive, non_interactive, more_archives, silent, **kwargs):  # pylint: disable=unused-argument
     """Utility function for `verdi import` to migrate archive
     Invoke click command `verdi export migrate`, passing in the archive,
     outputting the migrated archive in a temporary SandboxFolder.
@@ -107,6 +153,8 @@ def _migrate_archive(ctx, temp_folder, file_to_import, archive, non_interactive,
     :param file_to_import: Absolute path, including filename, of file to be migrated.
     :param archive: Filename of archive to be migrated, and later attempted imported.
     :param non_interactive: Whether or not the user should be asked for input for any reason.
+    :param more_archives: Whether or not there are more archives to be imported.
+    :param silent: Suppress console messages.
     :return: Absolute path to migrated archive within SandboxFolder.
     """
     from aiida.cmdline.commands.cmd_export import migrate
@@ -120,18 +168,19 @@ def _migrate_archive(ctx, temp_folder, file_to_import, archive, non_interactive,
     # Migration
     try:
         ctx.invoke(
-            migrate, input_file=file_to_import, output_file=temp_folder.get_abs_path(temp_out_file), silent=False
+            migrate, input_file=file_to_import, output_file=temp_folder.get_abs_path(temp_out_file), silent=silent
         )
-    except Exception:
-        echo.echo_error(
+    except Exception as exception:
+        _echo_error(
             'an exception occurred while migrating the archive {}.\n'
-            "Use 'verdi export migrate' to update this export file.".format(archive)
+            "Use 'verdi export migrate' to update this export file.".format(archive),
+            non_interactive=non_interactive,
+            more_archives=more_archives,
+            raised_exception=exception
         )
-        echo.echo(traceback.format_exc())
-        if not non_interactive:
-            click.confirm('do you want to continue?', abort=True)
     else:
-        echo.echo_success('archive migrated, proceeding with import')
+        # Success
+        echo.echo_info('proceeding with import')
 
         return temp_folder.get_abs_path(temp_out_file)
 
@@ -197,7 +246,6 @@ def cmd_import(
 
     The archive can be specified by its relative or absolute file path, or its HTTP URL.
     """
-
     from aiida.common.folders import SandboxFolder
     from aiida.tools.importexport.common.utils import get_valid_import_links
 
@@ -217,11 +265,13 @@ def cmd_import(
             try:
                 echo.echo_info('retrieving archive URLS from {}'.format(webpage))
                 urls = get_valid_import_links(webpage)
-            except Exception:
-                echo.echo_error('an exception occurred while trying to discover archives at URL {}'.format(webpage))
-                echo.echo(traceback.format_exc())
-                if not non_interactive:
-                    click.confirm('do you want to continue?', abort=True)
+            except Exception as exception:
+                _echo_error(
+                    'an exception occurred while trying to discover archives at URL {}'.format(webpage),
+                    non_interactive=non_interactive,
+                    more_archives=webpage != webpages[-1] or archives_file or archives_url,
+                    raised_exception=exception
+                )
             else:
                 echo.echo_success('{} archive URLs discovered and added'.format(len(urls)))
                 archives_url += urls
@@ -239,7 +289,8 @@ def cmd_import(
         'extras_mode_existing': ExtrasImportCode[extras_mode_existing].value,
         'extras_mode_new': extras_mode_new,
         'comment_mode': comment_mode,
-        'non_interactive': non_interactive
+        'non_interactive': non_interactive,
+        'silent': False,
     }
 
     # Import local archives
@@ -250,6 +301,7 @@ def cmd_import(
         # Initialization
         import_opts['archive'] = archive
         import_opts['file_to_import'] = import_opts['archive']
+        import_opts['more_archives'] = archive != archives_file[-1] or archives_url
 
         # First attempt to import archive
         migrate_archive = _try_import(migration_performed=False, **import_opts)
@@ -265,13 +317,14 @@ def cmd_import(
 
         # Initialization
         import_opts['archive'] = archive
+        import_opts['more_archives'] = archive != archives_url[-1]
 
         echo.echo_info('downloading archive {}'.format(archive))
 
         try:
             response = urllib.request.urlopen(archive)
         except Exception as exception:
-            echo.echo_warning('downloading archive {} failed: {}'.format(archive, exception))
+            _echo_error('downloading archive {} failed'.format(archive), raised_exception=exception, **import_opts)
 
         with SandboxFolder() as temp_folder:
             temp_file = 'importfile.tar.gz'

aiida/cmdline/commands/cmd_restapi.py

@@ -16,7 +16,7 @@
 import click
 
 from aiida.cmdline.commands.cmd_verdi import verdi
-from aiida.cmdline.params.options import HOSTNAME, PORT
+from aiida.cmdline.params.options import HOSTNAME, PORT, DEBUG
 from aiida.restapi.common import config
 
 
@@ -30,7 +30,7 @@
     default=config.CLI_DEFAULTS['CONFIG_DIR'],
     help='Path to the configuration directory'
 )
-@click.option('--debug', 'debug', is_flag=True, default=config.APP_CONFIG['DEBUG'], help='Enable debugging')
+@DEBUG(default=config.APP_CONFIG['DEBUG'])
 @click.option(
     '--wsgi-profile',
     is_flag=True,

aiida/cmdline/params/options/__init__.py

@@ -30,7 +30,7 @@
     'DESCRIPTION', 'INPUT_PLUGIN', 'CALC_JOB_STATE', 'PROCESS_STATE', 'PROCESS_LABEL', 'TYPE_STRING', 'EXIT_STATUS',
     'FAILED', 'LIMIT', 'PROJECT', 'ORDER_BY', 'PAST_DAYS', 'OLDER_THAN', 'ALL', 'ALL_STATES', 'ALL_USERS',
     'GROUP_CLEAR', 'RAW', 'HOSTNAME', 'TRANSPORT', 'SCHEDULER', 'USER', 'PORT', 'FREQUENCY', 'VERBOSE', 'TIMEOUT',
-    'FORMULA_MODE', 'TRAJECTORY_INDEX', 'WITH_ELEMENTS', 'WITH_ELEMENTS_EXCLUSIVE'
+    'FORMULA_MODE', 'TRAJECTORY_INDEX', 'WITH_ELEMENTS', 'WITH_ELEMENTS_EXCLUSIVE', 'DEBUG'
 )
 
 TRAVERSAL_RULE_HELP_STRING = {
@@ -522,3 +522,7 @@ def decorator(command):
 DICT_KEYS = OverridableOption(
     '-k', '--keys', type=click.STRING, cls=MultipleValueOption, help='Filter the output by one or more keys.'
 )
+
+DEBUG = OverridableOption(
+    '--debug', is_flag=True, default=False, help='Show debug messages. Mostly relevant for developers.', hidden=True
+)

aiida/common/folders.py

@@ -86,13 +86,13 @@ def get_subfolder(self, subfolder, create=False, reset_limit=False):
         Return a Folder object pointing to a subfolder.
 
         :param subfolder: a string with the relative path of the subfolder,
-                relative to the absolute path of this object. Note that
-                this may also contain '..' parts,
-                as far as this does not go beyond the folder_limit.
+            relative to the absolute path of this object. Note that
+            this may also contain '..' parts,
+            as far as this does not go beyond the folder_limit.
         :param create: if True, the new subfolder is created, if it does not exist.
         :param reset_limit: when doing ``b = a.get_subfolder('xxx', reset_limit=False)``,
-                the limit of b will be the same limit of a.
-                if True, the limit will be set to the boundaries of folder b.
+            the limit of b will be the same limit of a.
+            if True, the limit will be set to the boundaries of folder b.
 
         :Returns: a Folder object pointing to the subfolder.
         """
@@ -114,18 +114,16 @@ def get_subfolder(self, subfolder, create=False, reset_limit=False):
         return new_folder
 
     def get_content_list(self, pattern='*', only_paths=True):
-        """
-        Return a list of files (and subfolders) in the folder,
-        matching a given pattern.
+        """Return a list of files (and subfolders) in the folder, matching a given pattern.
 
         Example: If you want to exclude files starting with a dot, you can
         call this method with ``pattern='[!.]*'``
 
         :param pattern: a pattern for the file/folder names, using Unix filename
-                pattern matching (see Python standard module fnmatch).
-                By default, pattern is '*', matching all files and folders.
+            pattern matching (see Python standard module fnmatch).
+            By default, pattern is '*', matching all files and folders.
         :param only_paths: if False (default), return pairs (name, is_file).
-                if True, return only a flat list.
+            if True, return only a flat list.
 
         :Returns:
             a list of tuples of two elements, the first is the file name and
@@ -140,8 +138,7 @@ def get_content_list(self, pattern='*', only_paths=True):
         return [(fname, not os.path.isdir(os.path.join(self.abspath, fname))) for fname in file_list]
 
     def create_symlink(self, src, name):
-        """
-        Create a symlink inside the folder to the location 'src'.
+        """Create a symlink inside the folder to the location 'src'.
 
         :param src: the location to which the symlink must point. Can be
                 either a relative or an absolute path. Should, however,
@@ -155,8 +152,7 @@ def create_symlink(self, src, name):
         # For symlinks, permissions should not be set
 
     def insert_path(self, src, dest_name=None, overwrite=True):
-        """
-        Copy a file to the folder.
+        """Copy a file to the folder.
 
         :param src: the source filename to copy
         :param dest_name: if None, the same basename of src is used. Otherwise,
@@ -236,8 +232,7 @@ def create_file_from_filelike(self, filelike, filename, mode='wb', encoding=None
         return filepath
 
     def remove_path(self, filename):
-        """
-        Remove a file or folder from the folder.
+        """Remove a file or folder from the folder.
 
         :param filename: the relative path name to remove
         """
@@ -251,8 +246,7 @@ def remove_path(self, filename):
             os.remove(dest_abs_path)
 
     def get_abs_path(self, relpath, check_existence=False):
-        """
-        Return an absolute path for a file or folder in this folder.
+        """Return an absolute path for a file or folder in this folder.
 
         The advantage of using this method is that it checks that filename
         is a valid filename within this folder,
@@ -352,24 +346,20 @@ def create(self):
             os.makedirs(self.abspath, mode=self.mode_dir)
 
     def replace_with_folder(self, srcdir, move=False, overwrite=False):
-        """
-        This routine copies or moves the source folder 'srcdir' to the local
-        folder pointed by this Folder object.
+        """This routine copies or moves the source folder 'srcdir' to the local folder pointed to by this Folder.
 
-        :param srcdir: the source folder on the disk; this must be a string with
-                an absolute path
-        :param move: if True, the srcdir is moved to the repository. Otherwise, it
-                is only copied.
+        :param srcdir: the source folder on the disk; this must be an absolute path
+        :type srcdir: str
+        :param move: if True, the srcdir is moved to the repository. Otherwise, it is only copied.
+        :type move: bool
         :param overwrite: if True, the folder will be erased first.
-                if False, a IOError is raised if the folder already exists.
-                Whatever the value of this flag, parent directories will be
-                created, if needed.
+            if False, an IOError is raised if the folder already exists.
+            Whatever the value of this flag, parent directories will be created, if needed.
+        :type overwrite: bool
 
-        :Raises:
-            OSError or IOError: in case of problems accessing or writing
-            the files.
-        :Raises:
-            ValueError: if the section is not recognized.
+        :raises IOError: in case of problems accessing or writing the files.
+        :raises OSError: in case of problems accessing or writing the files (from ``shutil`` module).
+        :raises ValueError: if the section is not recognized.
         """
         if not os.path.isabs(srcdir):
             raise ValueError('srcdir must be an absolute path')
@@ -390,13 +380,11 @@ def replace_with_folder(self, srcdir, move=False, overwrite=False):
 
         # Set the mode also for the current dir, recursively
         for dirpath, _, filenames in os.walk(self.abspath, followlinks=False):
-            # dirpath should already be absolute, because I am passing
-            # an absolute path to os.walk
+            # dirpath should already be absolute, because I am passing an absolute path to os.walk
             os.chmod(dirpath, self.mode_dir)
             for filename in filenames:
-                # do not change permissions of symlinks (this would
-                # actually change permissions of the linked file/dir)
-                # Toc check whether this is a big speed loss
+                # do not change permissions of symlinks (this would actually change permissions of the linked file/dir)
+                # TODO check whether this is a big speed loss  # pylint: disable=fixme
                 full_file_path = os.path.join(dirpath, filename)
                 if not os.path.islink(full_file_path):
                     os.chmod(full_file_path, self.mode_file)