From 06205886204c142f771dab37f1a78f3bf0ba7251 Mon Sep 17 00:00:00 2001 From: Daniel Hollas Date: Sun, 15 Oct 2023 12:29:43 +0100 Subject: [PATCH] CLI: Do not load config in defaults and callbacks during tab-completion (#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`. --- aiida/cmdline/params/options/__init__.py | 2 + aiida/cmdline/params/options/callable.py | 34 +++++++++++++++ aiida/cmdline/params/options/interactive.py | 2 +- aiida/cmdline/params/options/main.py | 6 ++- tests/cmdline/params/options/test_callable.py | 42 +++++++++++++++++++ 5 files changed, 84 insertions(+), 2 deletions(-) create mode 100644 aiida/cmdline/params/options/callable.py create mode 100644 tests/cmdline/params/options/test_callable.py diff --git a/aiida/cmdline/params/options/__init__.py b/aiida/cmdline/params/options/__init__.py index b509d4e0ba..e5dd70f91b 100644 --- a/aiida/cmdline/params/options/__init__.py +++ b/aiida/cmdline/params/options/__init__.py @@ -14,6 +14,7 @@ # yapf: disable # pylint: disable=wildcard-import +from .callable import * from .config import * from .main import * from .multivalue import * @@ -39,6 +40,7 @@ 'COMPUTER', 'COMPUTERS', 'CONFIG_FILE', + 'CallableDefaultOption', 'ConfigFileOption', 'DATA', 'DATUM', diff --git a/aiida/cmdline/params/options/callable.py b/aiida/cmdline/params/options/callable.py new file mode 100644 index 0000000000..4fde0e96b6 --- /dev/null +++ b/aiida/cmdline/params/options/callable.py @@ -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) diff --git a/aiida/cmdline/params/options/interactive.py b/aiida/cmdline/params/options/interactive.py index 987016d598..f2f5489b56 100644 --- a/aiida/cmdline/params/options/interactive.py +++ b/aiida/cmdline/params/options/interactive.py @@ -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) diff --git a/aiida/cmdline/params/options/main.py b/aiida/cmdline/params/options/main.py index 60bc1a62de..ec431c3cb7 100644 --- a/aiida/cmdline/params/options/main.py +++ b/aiida/cmdline/params/options/main.py @@ -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 @@ -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: @@ -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.' ) diff --git a/tests/cmdline/params/options/test_callable.py b/tests/cmdline/params/options/test_callable.py new file mode 100644 index 0000000000..775b377395 --- /dev/null +++ b/tests/cmdline/params/options/test_callable.py @@ -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