Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Verdi: Load config / profile lazily to speed up tab-completion #6140

Merged
merged 2 commits into from
Oct 11, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
38 changes: 31 additions & 7 deletions aiida/cmdline/groups/verdi.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import base64
import difflib
import gzip
import typing as t

import click

Expand Down Expand Up @@ -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):
Expand Down
8 changes: 4 additions & 4 deletions aiida/cmdline/params/options/commands/setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)
Expand Down
3 changes: 2 additions & 1 deletion tests/cmdline/commands/test_setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test started failing with the changes introduced here. I am unsure if that means that I broke something, or whether the test worked by accident.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This actually looks like a regression because you made the default for this option lazy with a functools partial. I don't think the partial is properly called, so the value that is passed for the option is the partial and not the actual config option value

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Okay, after going through the guts of click and its modification in aiida, I've determined that the default value seems to be evaluating correctly, i.e. the callback is called, but it returns an empty tuple. I think the test worked previously because the default was evaluated before the actual test, and hence the get_config('autofill.user.email') was called in the local context, not in test context, where this value seems to be unavailable.

@sphuber could you maybe try if you can make the test fail locally on main branch in a fresh aiida install without any profile? That would I think confirm my hypothesis.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You are right. The problem is that the cmd_setup module is imported at the top of the test file. This will evaluate the setup command and with it the defaults. So even if I add the empty_config fixture to the test_setup_profile_uuid test, by the time that fixture gets invoked to provide an empty config for the test, the defaults of the setup command will already have been set. So by making the defaults lazily evaluated, this would indeed explain the observed behavior.

]

self.cli_runner(cmd_setup.setup, options, use_subprocess=False)
Expand Down