Add typing to aiida.tools.graph module (#7036)

Author: danielhollas

.pre-commit-config.yaml

@@ -141,10 +141,6 @@ repos:
         src/aiida/tools/data/orbital/orbital.py|
         src/aiida/tools/data/orbital/realhydrogen.py|
         src/aiida/tools/dbimporters/plugins/.*|
-        src/aiida/tools/graph/age_entities.py|
-        src/aiida/tools/graph/age_rules.py|
-        src/aiida/tools/graph/deletions.py|
-        src/aiida/tools/graph/graph_traversers.py|
         src/aiida/transports/cli.py|
         src/aiida/transports/plugins/local.py|
         src/aiida/transports/plugins/ssh.py|

docs/source/nitpick-exceptions

@@ -123,6 +123,8 @@ py:class FC
 py:class P
 py:class N
 py:class T
+py:class _ContainerTypes
+py:class _NodeOrGroupCls
 py:class aiida.cmdline.params.types.choice.T
 py:class aiida.common.lang.T
 py:class aiida.engine.processes.functions.N
@@ -162,6 +164,7 @@ py:class requests.Response
 py:class concurrent.futures._base.TimeoutError
 py:class concurrent.futures._base.Future
 
+py:class BackupManager
 py:class disk_objectstore.utils.LazyOpener
 py:class disk_objectstore.backup_utils.BackupManager
 

pyproject.toml

@@ -353,6 +353,7 @@ disallow_subclassing_any = true
 disallow_untyped_calls = true
 disallow_untyped_defs = true
 module = [
+  'aiida.tools.graph.*',
   'aiida.cmdline.params.*',
   'aiida.cmdline.groups.*',
   'aiida.tools.query.*'

src/aiida/cmdline/params/types/path.py

@@ -16,17 +16,11 @@
 
 import click
 
-if t.TYPE_CHECKING:
-    try:
-        from typing import TypeAlias
-    except ImportError:
-        from typing_extensions import TypeAlias
-
 __all__ = ('AbsolutePathParamType', 'FileOrUrl', 'PathOrUrl')
 
 URL_TIMEOUT_SECONDS = 10
 
-PathType: TypeAlias = 'str | bytes | os.PathLike[str]'
+PathType = t.Union[str, bytes, os.PathLike[str]]
 
 
 def check_timeout_seconds(timeout_seconds: float) -> int:

src/aiida/common/lang.py

@@ -13,7 +13,7 @@
 from typing import Any, Callable, Generic, TypeVar
 
 
-def isidentifier(identifier):
+def isidentifier(identifier: str) -> bool:
     """Return whether the given string is a valid python identifier.
 
     :return: boolean, True if identifier is valid, False otherwise
@@ -26,7 +26,7 @@ def isidentifier(identifier):
 T = TypeVar('T')
 
 
-def type_check(what: T, of_type: Any, msg: 'str | None' = None, allow_none: bool = False) -> 'T | None':
+def type_check(what: T, of_type: Any, msg: 'str | None' = None, allow_none: bool = False) -> T:
     """Verify that object 'what' is of type 'of_type' and if not the case, raise a TypeError.
 
     :param what: the object to check
@@ -37,7 +37,7 @@ def type_check(what: T, of_type: Any, msg: 'str | None' = None, allow_none: bool
     :return: `what` or `None`
     """
     if allow_none and what is None:
-        return None
+        return what
 
     if not isinstance(what, of_type):
         if msg is None:
@@ -50,7 +50,7 @@ def type_check(what: T, of_type: Any, msg: 'str | None' = None, allow_none: bool
 MethodType = TypeVar('MethodType', bound=Callable[..., Any])
 
 
-def override_decorator(check=False) -> Callable[[MethodType], MethodType]:
+def override_decorator(check: bool = False) -> Callable[[MethodType], MethodType]:
     """Decorator to signal that a method from a base class is being overridden completely."""
 
     def wrap(func: MethodType) -> MethodType:
@@ -63,18 +63,17 @@ def wrap(func: MethodType) -> MethodType:
         if not args:
             raise RuntimeError('Can only use the override decorator on member functions')
 
-        if check:
+        if not check:
+            return func
 
-            @functools.wraps(func)
-            def wrapped_fn(self, *args, **kwargs):
-                try:
-                    getattr(super(), func.__name__)
-                except AttributeError:
-                    raise RuntimeError(f'Function {func} does not override a superclass method')
+        @functools.wraps(func)
+        def wrapped_fn(self, *args, **kwargs):
+            try:
+                getattr(super(), func.__name__)
+            except AttributeError:
+                raise RuntimeError(f'Function {func} does not override a superclass method')
 
-                return func(self, *args, **kwargs)
-        else:
-            wrapped_fn = func  # type: ignore[assignment]
+            return func(self, *args, **kwargs)
 
         return wrapped_fn  # type: ignore[return-value]
 
@@ -84,7 +83,6 @@ def wrapped_fn(self, *args, **kwargs):
 override = override_decorator(check=False)
 
 ReturnType = TypeVar('ReturnType')
-SelfType = TypeVar('SelfType')
 
 
 class classproperty(Generic[ReturnType]):  # noqa: N801
@@ -95,8 +93,8 @@ class classproperty(Generic[ReturnType]):  # noqa: N801
     instance as its first argument).
     """
 
-    def __init__(self, getter: Callable[[SelfType], ReturnType]) -> None:
+    def __init__(self, getter: Callable[[Any], ReturnType]) -> None:
         self.getter = getter
 
-    def __get__(self, instance: Any, owner: SelfType) -> ReturnType:
-        return self.getter(owner)  # type: ignore[arg-type]
+    def __get__(self, instance: Any, owner: type) -> ReturnType:
+        return self.getter(owner)

src/aiida/common/log.py

@@ -10,7 +10,7 @@
 
 from __future__ import annotations
 
-import collections
+import collections.abc
 import contextlib
 import enum
 import io

src/aiida/common/typing.py

@@ -6,19 +6,19 @@
 # For further information on the license, see the LICENSE.txt file        #
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
-"""Module to define commonly used data structures."""
+"""Module to define commonly used types."""
 
 from __future__ import annotations
 
 import pathlib
+import sys
 from typing import Union
 
-try:
+if sys.version_info >= (3, 11):
     from typing import Self
-except ImportError:
+else:
     from typing_extensions import Self
 
 __all__ = ('FilePath', 'Self')
 
-
 FilePath = Union[str, pathlib.PurePath]

src/aiida/common/utils.py

@@ -18,20 +18,13 @@
 import sys
 from collections.abc import Iterable, Iterator
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, overload
+from typing import Any, Callable, TypeVar, overload
 from uuid import UUID
 
 from aiida.common.typing import Self
 
 from .lang import classproperty
 
-if TYPE_CHECKING:
-    # TypeAlias added in Python 3.10
-    try:
-        from typing import TypeAlias
-    except ImportError:
-        from typing_extensions import TypeAlias
-
 T = TypeVar('T')
 R = TypeVar('R')
 
@@ -421,7 +414,7 @@ def prettify(self, label: str) -> str:
         return self._prettifier_f(label)
 
 
-_Labels: TypeAlias = list[tuple[float, str]]
+_Labels = list[tuple[float, str]]
 
 
 def prettify_labels(labels: _Labels, format: str | None = None) -> _Labels:

src/aiida/orm/groups.py

@@ -12,11 +12,12 @@
 import warnings
 from functools import cached_property
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Sequence, Tuple, Type, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Any, ClassVar, Dict, Optional, Sequence, Tuple, Type, Union, cast
 
 from aiida.common import exceptions
 from aiida.common.lang import classproperty, type_check
 from aiida.common.pydantic import MetadataField
+from aiida.common.typing import Self
 from aiida.common.warnings import warn_deprecation
 from aiida.manage import get_manager
 
@@ -31,8 +32,6 @@
 
 __all__ = ('AutoGroup', 'Group', 'ImportGroup', 'UpfFamily')
 
-SelfType = TypeVar('SelfType', bound='Group')
-
 
 def load_group_class(type_string: str) -> Type['Group']:
     """Load the sub class of `Group` that corresponds to the given `type_string`.
@@ -203,7 +202,7 @@ def __repr__(self) -> str:
     def __str__(self) -> str:
         return f'{self.__class__.__name__}<{self.label}>'
 
-    def store(self: SelfType) -> SelfType:
+    def store(self) -> Self:
         """Verify that the group is allowed to be stored, which is the case along as `type_string` is set."""
         if self._type_string is None:
             raise exceptions.StoringNotAllowed('`type_string` is `None` so the group cannot be stored.')

src/aiida/orm/implementation/storage_backend.py

@@ -11,9 +11,12 @@
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING, Any, ContextManager, List, Optional, Sequence, TypeVar, Union
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Any, ContextManager, List, Optional, TypeVar, Union
 
 if TYPE_CHECKING:
+    from disk_objectstore.backup_utils import BackupManager
+
     from aiida.manage.configuration.profile import Profile
     from aiida.orm.autogroup import AutogroupManager
     from aiida.orm.entities import EntityTypes
@@ -129,7 +132,7 @@ def version(self) -> str:
         return version
 
     @abc.abstractmethod
-    def close(self):
+    def close(self) -> None:
         """Close the storage access."""
 
     @property
@@ -253,12 +256,12 @@ def delete(self) -> None:
         raise NotImplementedError()
 
     @abc.abstractmethod
-    def delete_nodes_and_connections(self, pks_to_delete: Sequence[int]):
+    def delete_nodes_and_connections(self, pks_to_delete: Iterable[int]) -> None:
         """Delete all nodes corresponding to pks in the input and any links to/from them.
 
         This method is intended to be used within a transaction context.
 
-        :param pks_to_delete: a sequence of node pks to delete
+        :param pks_to_delete: an iterable of node pks to delete
 
         :raises: ``AssertionError`` if a transaction is not active
         """
@@ -269,7 +272,7 @@ def get_repository(self) -> 'AbstractRepositoryBackend':
 
     @abc.abstractmethod
     def set_global_variable(
-        self, key: str, value: Union[None, str, int, float], description: Optional[str] = None, overwrite=True
+        self, key: str, value: Union[None, str, int, float], description: Optional[str] = None, overwrite: bool = True
     ) -> None:
         """Set a global variable in the storage.
 
@@ -291,7 +294,7 @@ def get_global_variable(self, key: str) -> Union[None, str, int, float]:
         """
 
     @abc.abstractmethod
-    def maintain(self, full: bool = False, dry_run: bool = False, **kwargs) -> None:
+    def maintain(self, full: bool = False, dry_run: bool = False, **kwargs: Any) -> None:
         """Perform maintenance tasks on the storage.
 
         If `full == True`, then this method may attempt to block the profile associated with the
@@ -309,10 +312,10 @@ def _backup(
         self,
         dest: str,
         keep: Optional[int] = None,
-    ):
+    ) -> None:
         raise NotImplementedError
 
-    def _write_backup_config(self, backup_manager):
+    def _write_backup_config(self, backup_manager: BackupManager) -> None:
         import pathlib
         import tempfile
 
@@ -338,7 +341,7 @@ def _write_backup_config(self, backup_manager):
         except (exceptions.MissingConfigurationError, exceptions.ConfigurationError) as exc:
             raise exceptions.StorageBackupError('AiiDA config.json not found!') from exc
 
-    def _validate_or_init_backup_folder(self, dest, keep):
+    def _validate_or_init_backup_folder(self, dest: str, keep: int | None) -> BackupManager:
         import json
         import tempfile
 
@@ -396,7 +399,7 @@ def backup(
         self,
         dest: str,
         keep: Optional[int] = None,
-    ):
+    ) -> None:
         """Create a backup of the storage contents.
 
         :param dest: The path to the destination folder.
@@ -440,15 +443,15 @@ def backup(
         STORAGE_LOGGER.report(f'Overwriting the `{DEFAULT_CONFIG_FILE_NAME} file.')
         self._write_backup_config(backup_manager)
 
-    def get_info(self, detailed: bool = False) -> dict:
+    def get_info(self, detailed: bool = False) -> dict[str, Any]:
         """Return general information on the storage.
 
         :param detailed: flag to request more detailed information about the content of the storage.
         :returns: a nested dict with the relevant information.
         """
         return {'entities': self.get_orm_entities(detailed=detailed)}
 
-    def get_orm_entities(self, detailed: bool = False) -> dict:
+    def get_orm_entities(self, detailed: bool = False) -> dict[str, Any]:
         """Return a mapping with an overview of the storage contents regarding ORM entities.
 
         :param detailed: flag to request more detailed information about the content of the storage.