Refactor: Delay import of heavy packages to speed up import time (#6116)

The importing of packages `urllib`, `yaml` and `pgsu` are moved from
top-level to inside the scopes where they are needed. This significantly
improves the load time of the `aiida` package and its subpackages.

The import of `Transport` in the `aiida.engine` package also slows down
imports, but it is only used for type checking, so its import is placed
inside the `if TYPE_CHECKING` guard.

Finally, the `DEFAULT_DBINFO`, `Postgres` and `PostgresConnectionMode`
objects of the `aiida.manage.external.postgres` package are no longer
exposed on the top-level as this also slows down imports. This is a
breaking change technically, but these resources should not be used by
downstream packages.

aiida/cmdline/commands/cmd_archive.py

@@ -14,7 +14,6 @@
 from pathlib import Path
 import traceback
 from typing import List, Tuple
-import urllib.request
 
 import click
 from click_spinner import spinner
@@ -432,6 +431,8 @@ def _import_archive_and_migrate(
     :param try_migration: whether to try a migration if the import raises `IncompatibleStorageSchema`
 
     """
+    import urllib.request
+
     from aiida.common.folders import SandboxFolder
     from aiida.tools.archive.abstract import get_format
     from aiida.tools.archive.imports import import_archive as _import_archive

aiida/cmdline/commands/cmd_code.py

@@ -13,7 +13,6 @@
 
 import click
 import tabulate
-import yaml
 
 from aiida.cmdline.commands.cmd_verdi import verdi
 from aiida.cmdline.groups.dynamic import DynamicEntryPointCommandGroup
@@ -242,6 +241,8 @@ def show(code):
 @with_dbenv()
 def export(code, output_file):
     """Export code to a yaml file."""
+    import yaml
+
     code_data = {}
 
     for key in code.get_cli_options().keys():

aiida/cmdline/commands/cmd_devel.py

@@ -66,8 +66,18 @@ def devel_check_undesired_imports():
     loaded_modules = 0
 
     for modulename in [
-        'asyncio', 'requests', 'plumpy', 'disk_objectstore', 'paramiko', 'seekpath', 'CifFile', 'ase', 'pymatgen',
-        'spglib', 'pymysql'
+        'asyncio',
+        'requests',
+        'plumpy',
+        'disk_objectstore',
+        'paramiko',
+        'seekpath',
+        'CifFile',
+        'ase',
+        'pymatgen',
+        'spglib',
+        'pymysql',
+        'yaml',
     ]:
         if modulename in sys.modules:
             echo.echo_warning(f'Detected loaded module "{modulename}"')

aiida/cmdline/commands/cmd_rabbitmq.py

@@ -18,7 +18,6 @@
 import click
 import tabulate
 import wrapt
-import yaml
 
 from aiida.cmdline.commands.cmd_devel import verdi_devel
 from aiida.cmdline.params import arguments, options
@@ -146,6 +145,8 @@ def with_manager(wrapped, _, args, kwargs):
 @with_manager
 def cmd_server_properties(manager):
     """List the server properties."""
+    import yaml
+
     data = {}
     for key, value in manager.get_communicator().server_properties.items():
         data[key] = value.decode('utf-8') if isinstance(value, bytes) else value

aiida/cmdline/params/options/config.py

@@ -24,7 +24,6 @@
 import typing as t
 
 import click
-import yaml
 
 from .overridable import OverridableOption
 
@@ -33,6 +32,8 @@
 
 def yaml_config_file_provider(handle, cmd_name):  # pylint: disable=unused-argument
     """Read yaml config file from file handle."""
+    import yaml
+
     return yaml.safe_load(handle)
 
 

aiida/cmdline/params/options/main.py

@@ -9,9 +9,9 @@
 ###########################################################################
 """Module with pre-defined reusable commandline options that can be used as `click` decorators."""
 import click
-from pgsu import DEFAULT_DSN as DEFAULT_DBINFO  # pylint: disable=no-name-in-module
 
 from aiida.common.log import LOG_LEVELS, configure_logging
+from aiida.manage.external.postgres import DEFAULT_DBINFO
 from aiida.manage.external.rmq import BROKER_DEFAULTS
 
 from .. import types

aiida/cmdline/params/types/path.py

@@ -10,8 +10,6 @@
 """Click parameter types for paths."""
 import os
 from socket import timeout
-import urllib.error
-import urllib.request
 
 import click
 
@@ -88,6 +86,9 @@ def convert(self, value, param, ctx):
 
     def checks_url(self, url, param, ctx):
         """Check whether URL is reachable within timeout."""
+        import urllib.error
+        import urllib.request
+
         try:
             with urllib.request.urlopen(url, timeout=self.timeout_seconds):
                 pass
@@ -123,6 +124,9 @@ def convert(self, value, param, ctx):
 
     def get_url(self, url, param, ctx):
         """Retrieve file from URL."""
+        import urllib.error
+        import urllib.request
+
         try:
             return urllib.request.urlopen(url, timeout=self.timeout_seconds)  # pylint: disable=consider-using-with
         except (urllib.error.URLError, urllib.error.HTTPError, timeout):

aiida/cmdline/utils/echo.py

@@ -16,7 +16,6 @@
 from typing import Any, Optional
 
 import click
-import yaml
 
 CMDLINE_LOGGER = logging.getLogger('verdi')
 
@@ -225,11 +224,15 @@ def default_jsondump(data):
 
 def _format_yaml(dictionary, sort_keys=True):
     """Return a dictionary formatted as a string using the YAML format."""
+    import yaml
+
     return yaml.dump(dictionary, sort_keys=sort_keys)
 
 
 def _format_yaml_expanded(dictionary, sort_keys=True):
     """Return a dictionary formatted as a string using the expanded YAML format."""
+    import yaml
+
     return yaml.dump(dictionary, sort_keys=sort_keys, default_flow_style=False)
 
 

aiida/engine/daemon/execmanager.py

@@ -21,9 +21,9 @@
 import pathlib
 import shutil
 from tempfile import NamedTemporaryFile
-from typing import Any, List
 from typing import Mapping as MappingType
 from typing import Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, List
 
 from aiida.common import AIIDA_LOGGER, exceptions
 from aiida.common.datastructures import CalcInfo
@@ -35,7 +35,9 @@
 from aiida.orm.utils.log import get_dblogger_extra
 from aiida.repository.common import FileType
 from aiida.schedulers.datastructures import JobState
-from aiida.transports import Transport
+
+if TYPE_CHECKING:
+    from aiida.transports import Transport
 
 REMOTE_WORK_DIRECTORY_LOST_FOUND = 'lost+found'
 

aiida/engine/processes/builder.py

@@ -13,8 +13,6 @@
 from typing import TYPE_CHECKING, Any, Type
 from uuid import uuid4
 
-import yaml
-
 from aiida.engine.processes.ports import PortNamespace
 from aiida.orm import Dict, Node
 from aiida.orm.nodes.data.base import BaseType
@@ -245,6 +243,8 @@ def process_class(self) -> Type['Process']:
 
     def _repr_pretty_(self, p, _) -> str:  # pylint: disable=invalid-name
         """Pretty representation for in the IPython console and notebooks."""
+        import yaml
+
         return p.text(
             f'Process class: {self._process_class.__name__}\n'
             f'Inputs:\n{yaml.safe_dump(json.JSONDecoder().decode(PrettyEncoder().encode(self)))}'

aiida/engine/processes/calcjobs/monitors.py

@@ -13,7 +13,9 @@
 from aiida.common.log import AIIDA_LOGGER
 from aiida.orm import CalcJobNode, Dict
 from aiida.plugins import BaseFactory
-from aiida.transports import Transport
+
+if t.TYPE_CHECKING:
+    from aiida.transports import Transport
 
 LOGGER = AIIDA_LOGGER.getChild(__name__)
 

aiida/engine/transports.py

@@ -13,10 +13,12 @@
 import contextvars
 import logging
 import traceback
-from typing import Awaitable, Dict, Hashable, Iterator, Optional
+from typing import TYPE_CHECKING, Awaitable, Dict, Hashable, Iterator, Optional
 
 from aiida.orm import AuthInfo
-from aiida.transports import Transport
+
+if TYPE_CHECKING:
+    from aiida.transports import Transport
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -55,7 +57,7 @@ def loop(self) -> asyncio.AbstractEventLoop:
         return self._loop
 
     @contextlib.contextmanager
-    def request_transport(self, authinfo: AuthInfo) -> Iterator[Awaitable[Transport]]:
+    def request_transport(self, authinfo: AuthInfo) -> Iterator[Awaitable['Transport']]:
         """
         Request a transport from an authinfo.  Because the client is not allowed to
         request a transport immediately they will instead be given back a future

aiida/manage/__init__.py

@@ -35,13 +35,10 @@
     'CURRENT_CONFIG_VERSION',
     'Config',
     'ConfigValidationError',
-    'DEFAULT_DBINFO',
     'MIGRATIONS',
     'ManagementApiConnectionError',
     'OLDEST_COMPATIBLE_CONFIG_VERSION',
     'Option',
-    'Postgres',
-    'PostgresConnectionMode',
     'Profile',
     'RabbitmqManagementClient',
     'check_and_migrate_config',

aiida/manage/external/__init__.py

@@ -14,15 +14,11 @@
 # yapf: disable
 # pylint: disable=wildcard-import
 
-from .postgres import *
 from .rmq import *
 
 __all__ = (
     'BROKER_DEFAULTS',
-    'DEFAULT_DBINFO',
     'ManagementApiConnectionError',
-    'Postgres',
-    'PostgresConnectionMode',
     'RabbitmqManagementClient',
     'get_launch_queue_name',
     'get_message_exchange_name',

aiida/manage/external/postgres.py

@@ -23,8 +23,6 @@
 if TYPE_CHECKING:
     from aiida.manage.configuration import Profile
 
-__all__ = ('Postgres', 'PostgresConnectionMode', 'DEFAULT_DBINFO')
-
 # The last placeholder is for adding privileges of the user
 _CREATE_USER_COMMAND = 'CREATE USER "{}" WITH PASSWORD \'{}\' {}'
 _DROP_USER_COMMAND = 'DROP USER "{}"'

aiida/manage/external/rmq/client.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import typing as t
-from urllib.parse import quote
 
 from aiida.common.exceptions import AiidaException
 
@@ -49,6 +48,8 @@ def format_url(self, url: str, url_params: dict[str, str] | None = None) -> str:
             automatically inserted and should not be specified.
         :returns: The complete URL.
         """
+        from urllib.parse import quote
+
         url_params = url_params or {}
         url_params['virtual_host'] = self._virtual_host if self._virtual_host else '/'
         url_params = {key: quote(value, safe='') for key, value in url_params.items()}

aiida/manage/external/rmq/utils.py

@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 """Utilites for RabbitMQ."""
-from urllib.parse import urlencode, urlunparse
 
 from . import defaults
 
@@ -25,6 +24,8 @@ def get_rmq_url(protocol=None, username=None, password=None, host=None, port=Non
     :param kwargs: remaining keyword arguments that will be encoded as query parameters.
     :returns: the connection URL string.
     """
+    from urllib.parse import urlencode, urlunparse
+
     if 'heartbeat' not in kwargs:
         kwargs['heartbeat'] = defaults.BROKER_DEFAULTS.heartbeat
 

aiida/schedulers/scheduler.py

@@ -19,7 +19,9 @@
 from aiida.common.lang import classproperty
 from aiida.engine.processes.exit_code import ExitCode
 from aiida.schedulers.datastructures import JobInfo, JobResource, JobTemplate, JobTemplateCodeInfo
-from aiida.transports import Transport
+
+if t.TYPE_CHECKING:
+    from aiida.transports import Transport
 
 __all__ = ('Scheduler', 'SchedulerError', 'SchedulerParsingError')
 

docs/source/nitpick-exceptions

@@ -8,6 +8,7 @@ py:class callable
 py:class function
 py:class traceback
 py:class NoneType
+py:class MappingType
 py:class AbstractContextManager
 py:class BinaryIO
 py:class IO