Skip to content

Commit

Permalink
CLI: Do not load config in defaults and callbacks during tab-completi…
Browse files Browse the repository at this point in the history
…on (#6144)

The `get_default_profile` default of the `PROFILE` option and the
`set_log_level` callback of the `VERBOSITY` option both load the config.
Since defaults and callbacks are also evaluated during tab-completion
this was slowing down tab-completion significantly since loading the
config has a non-negligible cost.

The `set_log_level` callback is modified to explicitly check whether
we are are tab-completing, in which case `ctx.resilient_parsing` is set
to `True`. In this case, the functions now returns `None` and no longer
loads the config.

For `get_default_profile`, the `CallableDefaultOption` class is added
which allows the default to be made a callable, which will return `None`
if `ctx.resilient_parsing` is set to `True`.
  • Loading branch information
danielhollas authored Oct 15, 2023
1 parent a1f456d commit 0620588
Show file tree
Hide file tree
Showing 5 changed files with 84 additions and 2 deletions.
2 changes: 2 additions & 0 deletions aiida/cmdline/params/options/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
# yapf: disable
# pylint: disable=wildcard-import

from .callable import *
from .config import *
from .main import *
from .multivalue import *
Expand All @@ -39,6 +40,7 @@
'COMPUTER',
'COMPUTERS',
'CONFIG_FILE',
'CallableDefaultOption',
'ConfigFileOption',
'DATA',
'DATUM',
Expand Down
34 changes: 34 additions & 0 deletions aiida/cmdline/params/options/callable.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved. #
# This file is part of the AiiDA code. #
# #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.aiida.net #
###########################################################################
"""
.. py:module::callable
:synopsis: A monkey-patched subclass of click.Option that does not evaluate callable default during tab completion
"""

import typing as t

import click

__all__ = ('CallableDefaultOption',)


class CallableDefaultOption(click.Option):
"""A monkeypatch for click.Option that does not evaluate default callbacks during tab completion
This is a temporary solution until a proper fix is implemented in click, see:
https://github.com/pallets/click/issues/2614
"""

def get_default(self, ctx: click.Context, call: bool = True) -> t.Optional[t.Union[t.Any, t.Callable[[], t.Any]]]:
"""provides the functionality of :meth:`click.Option.get_default`,
but ensures we do not evaluate callable defaults when in tab-completion context."""
if ctx.resilient_parsing:
return None
return super().get_default(ctx=ctx, call=call)
2 changes: 1 addition & 1 deletion aiida/cmdline/params/options/interactive.py
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ def get_default(self, ctx: click.Context, call: bool = True) -> t.Optional[t.Uni
if self._contextual_default is not None:
default = self._contextual_default(ctx)
else:
default = super().get_default(ctx)
default = super().get_default(ctx, call=call)

try:
default = self.type.deconvert_default(default)
Expand Down
6 changes: 5 additions & 1 deletion aiida/cmdline/params/options/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

from .. import types
from ...utils import defaults, echo # pylint: disable=no-name-in-module
from .callable import CallableDefaultOption
from .config import ConfigFileOption
from .multivalue import MultipleValueOption
from .overridable import OverridableOption
Expand Down Expand Up @@ -110,8 +111,10 @@ def set_log_level(_ctx, _param, value):
log.CLI_ACTIVE = True

# If the value is ``None``, it means the option was not specified, but we still configure logging for the CLI
# However, we skip this when we are in a tab-completion context.
if value is None:
configure_logging()
if not _ctx.resilient_parsing:
configure_logging()
return None

try:
Expand Down Expand Up @@ -145,6 +148,7 @@ def set_log_level(_ctx, _param, value):
'profile',
type=types.ProfileParamType(),
default=defaults.get_default_profile,
cls=CallableDefaultOption,
help='Execute the command for this profile instead of the default profile.'
)

Expand Down
42 changes: 42 additions & 0 deletions tests/cmdline/params/options/test_callable.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,42 @@
# -*- coding: utf-8 -*-
###########################################################################
# Copyright (c), The AiiDA team. All rights reserved. #
# This file is part of the AiiDA code. #
# #
# The code is hosted on GitHub at https://github.com/aiidateam/aiida-core #
# For further information on the license, see the LICENSE.txt file #
# For further information please visit http://www.aiida.net #
###########################################################################
# pylint: disable=redefined-outer-name
"""Tests for the :mod:`aiida.cmdline.params.options.callable` module."""
from click.shell_completion import ShellComplete
import pytest

from aiida.cmdline.commands.cmd_verdi import verdi


def _get_completions(cli, args, incomplete):
comp = ShellComplete(cli, {}, cli.name, '_CLICK_COMPLETE')
return comp.get_completions(args, incomplete)


@pytest.fixture
def unload_config():
"""Temporarily unload the config by setting ``aiida.manage.configuration.CONFIG`` to ``None``."""
from aiida.manage import configuration

config = configuration.CONFIG
configuration.CONFIG = None
yield
configuration.CONFIG = config


@pytest.mark.usefixtures('unload_config')
def test_callable_default_resilient_parsing():
"""Test that tab-completion of ``verdi`` does not evaluate defaults that load the config, which is expensive."""
from aiida.manage import configuration

assert configuration.CONFIG is None
completions = [c.value for c in _get_completions(verdi, [], '')]
assert 'help' in completions
assert configuration.CONFIG is None

0 comments on commit 0620588

Please sign in to comment.