From 64111abfbb28ace90a759fcb88c1d76db2ba58e7 Mon Sep 17 00:00:00 2001 From: Daniel Hollas Date: Tue, 26 Sep 2023 23:20:23 +0100 Subject: [PATCH 1/2] CLI: Make loading of config lazy 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 | 35 +++++++++++++++---- .../cmdline/params/options/commands/setup.py | 8 ++--- tests/cmdline/commands/test_setup.py | 3 +- 3 files changed, 34 insertions(+), 12 deletions(-) diff --git a/aiida/cmdline/groups/verdi.py b/aiida/cmdline/groups/verdi.py index d0bd9ed49f..735da0fd7c 100644 --- a/aiida/cmdline/groups/verdi.py +++ b/aiida/cmdline/groups/verdi.py @@ -34,19 +34,40 @@ ) +class LazyConfigAttributeDict(AttributeDict): + """We want to avoid needlesly loading user config on every verdi invocation""" + + _LAZY_KEY = 'config' + + def __init__(self, ctx, dictionary=None): + super().__init__(dictionary) + # We need click context so we can stop in case we fail to parse the config + self.ctx = ctx + + def __getattr__(self, attr): + """Override of AttributeDict.__getattr__ for lazy loading the config key. + + Stops if config file cannot be parsed. + :raises AttributeError: if the attribute does not correspond to an existing key. + """ + 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) From d7f78c3e1c9ae38446bf53fa8b535192581066da Mon Sep 17 00:00:00 2001 From: Sebastiaan Huber Date: Wed, 11 Oct 2023 10:54:35 +0200 Subject: [PATCH 2/2] minor touchups --- aiida/cmdline/groups/verdi.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/aiida/cmdline/groups/verdi.py b/aiida/cmdline/groups/verdi.py index 735da0fd7c..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 @@ -35,20 +36,21 @@ class LazyConfigAttributeDict(AttributeDict): - """We want to avoid needlesly loading user config on every verdi invocation""" + """Subclass of ``AttributeDict`` that lazily calls :meth:`aiida.manage.configuration.get_config`.""" _LAZY_KEY = 'config' - def __init__(self, ctx, dictionary=None): + def __init__(self, ctx: click.Context, dictionary: dict[str, t.Any] | None = None): super().__init__(dictionary) - # We need click context so we can stop in case we fail to parse the config self.ctx = ctx - def __getattr__(self, attr): - """Override of AttributeDict.__getattr__ for lazy loading the config key. + def __getattr__(self, attr: str) -> t.Any: + """Override of ``AttributeDict.__getattr__`` for lazily loading the config key. - Stops if config file cannot be parsed. - :raises AttributeError: if the attribute does not correspond to an existing 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) @@ -58,6 +60,7 @@ def __getattr__(self, attr): self[self._LAZY_KEY] = get_config(create=True) except ConfigurationError as exception: self.ctx.fail(str(exception)) + return self[self._LAZY_KEY]