CLI: Speed up tab-completion by lazily importing Config (#6275)

The `Config` class was recently changed to use `pydantic` which adds
a non-negligible load time to `verdi`, which affects the responsiveness
of tab-completion. To restore this, the `Config` class cannot be imported
upon loading `verdi` and so its exposure in the parent module had to be
removed. Instead of importing from `aiida.manage.configuration` it now
has to be imported from the full path `aiida.manage.configuration.config`.

src/aiida/cmdline/commands/cmd_config.py

@@ -38,7 +38,8 @@ def verdi_config_list(ctx, prefix, description: bool):
     """
     from tabulate import tabulate
 
-    from aiida.manage.configuration import Config, Profile
+    from aiida.manage.configuration import Profile
+    from aiida.manage.configuration.config import Config
 
     config: Config = ctx.obj.config
     profile: Profile | None = ctx.obj.get('profile', None)
@@ -79,7 +80,8 @@ def _join(val):
 @click.pass_context
 def verdi_config_show(ctx, option):
     """Show details of an AiiDA option for the current profile."""
-    from aiida.manage.configuration import Config, Profile
+    from aiida.manage.configuration import Profile
+    from aiida.manage.configuration.config import Config
 
     config: Config = ctx.obj.config
     profile: Profile | None = ctx.obj.profile
@@ -125,7 +127,10 @@ def verdi_config_set(ctx, option, value, globally, append, remove):
     import typing
 
     from aiida.common.exceptions import ConfigurationError
-    from aiida.manage.configuration import Config, Profile
+
+    if typing.TYPE_CHECKING:
+        from aiida.manage.configuration import Profile
+        from aiida.manage.configuration.config import Config
 
     if append and remove:
         echo.echo_critical('Cannot flag both append and remove')
@@ -170,7 +175,8 @@ def verdi_config_set(ctx, option, value, globally, append, remove):
 @click.pass_context
 def verdi_config_unset(ctx, option, globally):
     """Unset an AiiDA option."""
-    from aiida.manage.configuration import Config, Profile
+    from aiida.manage.configuration import Profile
+    from aiida.manage.configuration.config import Config
 
     config: Config = ctx.obj.config
     profile: Profile | None = ctx.obj.profile

src/aiida/manage/__init__.py

@@ -30,7 +30,6 @@
 __all__ = (
     'BROKER_DEFAULTS',
     'CURRENT_CONFIG_VERSION',
-    'Config',
     'MIGRATIONS',
     'ManagementApiConnectionError',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',

src/aiida/manage/configuration/__init__.py

@@ -13,14 +13,12 @@
 
 # fmt: off
 
-from .config import *
 from .migrations import *
 from .options import *
 from .profile import *
 
 __all__ = (
     'CURRENT_CONFIG_VERSION',
-    'Config',
     'MIGRATIONS',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',
     'Option',
@@ -58,7 +56,7 @@
 from aiida.common.warnings import AiidaDeprecationWarning
 
 if TYPE_CHECKING:
-    from aiida.manage.configuration import Config, Profile
+    from .config import Config
 
 # global variables for aiida
 CONFIG: Optional['Config'] = None
@@ -197,7 +195,7 @@ def profile_context(profile: Optional[str] = None, allow_switch=False) -> 'Profi
 
 
 def create_profile(
-    config: Config,
+    config: 'Config',
     storage_cls,
     *,
     name: str,

src/aiida/manage/configuration/config.py

@@ -37,8 +37,6 @@
 from .options import Option, get_option, get_option_names, parse_option
 from .profile import Profile
 
-__all__ = ('Config',)
-
 if TYPE_CHECKING:
     from aiida.orm.implementation.storage_backend import StorageBackend
 

src/aiida/manage/tests/pytest_fixtures.py

@@ -42,10 +42,13 @@
 from aiida.common.warnings import warn_deprecation
 from aiida.engine import Process, ProcessBuilder, submit
 from aiida.engine.daemon.client import DaemonClient, DaemonNotRunningException, DaemonTimeoutException
-from aiida.manage import Config, Profile, get_manager, get_profile
+from aiida.manage import Profile, get_manager, get_profile
 from aiida.manage.manager import Manager
 from aiida.orm import Computer, ProcessNode, User
 
+if t.TYPE_CHECKING:
+    from aiida.manage.configuration.config import Config
+
 
 def recursive_merge(left: dict[t.Any, t.Any], right: dict[t.Any, t.Any]) -> None:
     """Recursively merge the ``right`` dictionary into the ``left`` dictionary.

tests/cmdline/params/options/test_callable.py

@@ -7,10 +7,14 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """Tests for the :mod:`aiida.cmdline.params.options.callable` module."""
+import sys
+
 import pytest
 from aiida.cmdline.commands.cmd_verdi import verdi
 from click.shell_completion import ShellComplete
 
+SLOW_IMPORTS = ['pydantic']
+
 
 def _get_completions(cli, args, incomplete):
     comp = ShellComplete(cli, {}, cli.name, '_CLICK_COMPLETE')
@@ -28,6 +32,13 @@ def unload_config():
     configuration.CONFIG = config
 
 
+@pytest.fixture
+def unimport_slow_imports():
+    """Remove modules in ``SLOW_IMPORTS`` from ``sys.modules``."""
+    for module in SLOW_IMPORTS:
+        del sys.modules[module]
+
+
 @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."""
@@ -37,3 +48,22 @@ def test_callable_default_resilient_parsing():
     completions = [c.value for c in _get_completions(verdi, [], '')]
     assert 'help' in completions
     assert configuration.CONFIG is None
+
+
+@pytest.mark.usefixtures('unload_config')
+@pytest.mark.usefixtures('unimport_slow_imports')
+def test_slow_imports_during_tab_completion():
+    """Check that verdi does not import certain python modules that would make tab-completion slow.
+
+    NOTE: This is analogous to `verdi devel check-undesired-imports`
+    """
+
+    # Let's double check that the undesired imports are not already loaded
+    for modulename in SLOW_IMPORTS:
+        assert modulename not in sys.modules, f'Module `{modulename}` was not properly unloaded'
+
+    completions = [c.value for c in _get_completions(verdi, [], '')]
+    assert 'help' in completions
+
+    for modulename in SLOW_IMPORTS:
+        assert modulename not in sys.modules, f'Detected loaded module {modulename} during tab completion'

tests/conftest.py

@@ -24,7 +24,10 @@
 import click
 import pytest
 from aiida import get_profile
-from aiida.manage.configuration import Config, Profile, get_config, load_profile
+from aiida.manage.configuration import Profile, get_config, load_profile
+
+if t.TYPE_CHECKING:
+    from aiida.manage.configuration.config import Config
 
 pytest_plugins = ['aiida.manage.tests.pytest_fixtures', 'sphinx.testing.fixtures']
 
@@ -174,8 +177,9 @@ def isolated_config(monkeypatch):
     Python process and so doesn't have access to the loaded config in memory in the process that is running the test.
     """
     from aiida.manage import configuration
+    from aiida.manage.configuration.config import Config
 
-    monkeypatch.setattr(configuration.Config, '_backup', lambda *args, **kwargs: None)
+    monkeypatch.setattr(Config, '_backup', lambda *args, **kwargs: None)
 
     current_config = configuration.CONFIG
     configuration.CONFIG = copy.deepcopy(current_config)

tests/manage/configuration/test_config.py

@@ -14,7 +14,8 @@
 
 import pytest
 from aiida.common import exceptions
-from aiida.manage.configuration import Config, Profile, settings
+from aiida.manage.configuration import Profile, settings
+from aiida.manage.configuration.config import Config
 from aiida.manage.configuration.migrations import CURRENT_CONFIG_VERSION, OLDEST_COMPATIBLE_CONFIG_VERSION
 from aiida.manage.configuration.options import get_option
 from aiida.orm.implementation.storage_backend import StorageBackend