From d533b7a540ab9d420acec1833bb7e23f50d8a7c1 Mon Sep 17 00:00:00 2001 From: Daniel Hollas Date: Wed, 11 Oct 2023 10:43:55 +0100 Subject: [PATCH] CLI: Make loading of config lazy for improved responsiveness (#6140) The `VerdiContext` class, which provides the custom context of the `verdi` commands, loads the configuration. This has a non-negligible cost and so slows down the responsiveness of the CLI. This is especially noticeable during tab-completion. The `obj` custom object of the `VerdiContext` is replaced with a subclass of `AttributeDict` that lazily populates the `config` key when it is called with the loaded `Config` class. In addition, the defaults of some options of the `verdi setup` command, which load a value from the config and so require the config, are turned into partials such that they also are lazily evaluated. These changes should give a reduction in load time of `verdi` of the order of ~50 ms. A test of `verdi setup` had to be updated to explicitly provide a value for the email. This is because now the default is evaluated lazily, i.e. when the command is actually called in the test. At this point, there is no value for this config option and so the default is empty. Before, the default would be evaluated as soon as `aiida.cmdline.commands.cmd_setup` was imported, at which point an existing config would still contain these values, binding them to the default, even if the config would be reset afterwards before the test. --- aiida/cmdline/groups/verdi.py | 38 +++++++++++++++---- .../cmdline/params/options/commands/setup.py | 8 ++-- tests/cmdline/commands/test_setup.py | 3 +- 3 files changed, 37 insertions(+), 12 deletions(-) diff --git a/aiida/cmdline/groups/verdi.py b/aiida/cmdline/groups/verdi.py index d0bd9ed49f..91fcb88ba1 100644 --- a/aiida/cmdline/groups/verdi.py +++ b/aiida/cmdline/groups/verdi.py @@ -5,6 +5,7 @@ import base64 import difflib import gzip +import typing as t import click @@ -34,19 +35,42 @@ ) +class LazyConfigAttributeDict(AttributeDict): + """Subclass of ``AttributeDict`` that lazily calls :meth:`aiida.manage.configuration.get_config`.""" + + _LAZY_KEY = 'config' + + def __init__(self, ctx: click.Context, dictionary: dict[str, t.Any] | None = None): + super().__init__(dictionary) + self.ctx = ctx + + def __getattr__(self, attr: str) -> t.Any: + """Override of ``AttributeDict.__getattr__`` for lazily loading the config key. + + :param attr: The attribute to return. + :returns: The value of the attribute. + :raises AttributeError: If the attribute does not correspond to an existing key. + :raises click.exceptions.UsageError: If loading of the configuration fails. + """ + if attr != self._LAZY_KEY: + return super().__getattr__(attr) + + if self._LAZY_KEY not in self: + try: + self[self._LAZY_KEY] = get_config(create=True) + except ConfigurationError as exception: + self.ctx.fail(str(exception)) + + return self[self._LAZY_KEY] + + class VerdiContext(click.Context): """Custom context implementation that defines the ``obj`` user object and adds the ``Config`` instance.""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - if self.obj is None: - self.obj = AttributeDict() - - try: - self.obj.config = get_config(create=True) - except ConfigurationError as exception: - self.fail(str(exception)) + self.obj = LazyConfigAttributeDict(self) class VerdiCommandGroup(click.Group): diff --git a/aiida/cmdline/params/options/commands/setup.py b/aiida/cmdline/params/options/commands/setup.py index 2838948863..b3218bdda5 100644 --- a/aiida/cmdline/params/options/commands/setup.py +++ b/aiida/cmdline/params/options/commands/setup.py @@ -176,28 +176,28 @@ def get_quicksetup_password(ctx, param, value): # pylint: disable=unused-argume SETUP_USER_EMAIL = options.USER_EMAIL.clone( prompt='Email Address (for sharing data)', - default=get_config_option('autofill.user.email'), + default=functools.partial(get_config_option, 'autofill.user.email'), required=True, cls=options.interactive.InteractiveOption ) SETUP_USER_FIRST_NAME = options.USER_FIRST_NAME.clone( prompt='First name', - default=get_config_option('autofill.user.first_name'), + default=functools.partial(get_config_option, 'autofill.user.first_name'), required=True, cls=options.interactive.InteractiveOption ) SETUP_USER_LAST_NAME = options.USER_LAST_NAME.clone( prompt='Last name', - default=get_config_option('autofill.user.last_name'), + default=functools.partial(get_config_option, 'autofill.user.last_name'), required=True, cls=options.interactive.InteractiveOption ) SETUP_USER_INSTITUTION = options.USER_INSTITUTION.clone( prompt='Institution', - default=get_config_option('autofill.user.institution'), + default=functools.partial(get_config_option, 'autofill.user.institution'), required=True, cls=options.interactive.InteractiveOption ) diff --git a/tests/cmdline/commands/test_setup.py b/tests/cmdline/commands/test_setup.py index 887c474b48..68c1fcd9b8 100644 --- a/tests/cmdline/commands/test_setup.py +++ b/tests/cmdline/commands/test_setup.py @@ -205,11 +205,12 @@ def test_setup_profile_uuid(self): postgres.create_db(db_user, db_name) profile_name = 'profile-copy' + user_email = 'some@email.com' profile_uuid = str(uuid.uuid4) options = [ '--non-interactive', '--profile', profile_name, '--profile-uuid', profile_uuid, '--db-backend', self.storage_backend_name, '--db-name', db_name, '--db-username', db_user, '--db-password', db_pass, - '--db-port', self.pg_test.dsn['port'] + '--db-port', self.pg_test.dsn['port'], '--email', user_email ] self.cli_runner(cmd_setup.setup, options, use_subprocess=False)