Add typing to aiida/common/{utils,extendeddicts}.py

.pre-commit-config.yaml

@@ -84,8 +84,6 @@ repos:
         src/aiida/calculations/diff_tutorial/calculations.py|
         src/aiida/calculations/templatereplacer.py|
         src/aiida/calculations/transfer.py|
-        src/aiida/common/extendeddicts.py|
-        src/aiida/common/utils.py|
         src/aiida/engine/daemon/execmanager.py|
         src/aiida/engine/processes/calcjobs/manager.py|
         src/aiida/engine/processes/calcjobs/monitors.py|
@@ -95,7 +93,6 @@ repos:
         src/aiida/manage/configuration/__init__.py|
         src/aiida/manage/configuration/config.py|
         src/aiida/manage/external/rmq/launcher.py|
-        src/aiida/orm/nodes/data/array/bands.py|
         src/aiida/orm/nodes/data/array/trajectory.py|
         src/aiida/orm/nodes/data/cif.py|
         src/aiida/orm/nodes/data/remote/base.py|

src/aiida/cmdline/groups/verdi.py

@@ -6,6 +6,7 @@
 import difflib
 import gzip
 import typing as t
+from collections.abc import Hashable
 
 import click
 
@@ -48,10 +49,10 @@ class LazyVerdiObjAttributeDict(AttributeDict):
     _KEY_PROFILE = 'profile'
 
     def __init__(self, ctx: click.Context, dictionary: dict[str, t.Any] | None = None):
-        super().__init__(dictionary)  # type: ignore[no-untyped-call]
+        super().__init__(dictionary)
         self.ctx = ctx
 
-    def __getattr__(self, attr: str) -> t.Any:
+    def __getattr__(self, attr: Hashable) -> t.Any:
         """Override of ``AttributeDict.__getattr__`` to lazily initialize the ``config`` and ``profile`` attributes.
 
         :param attr: The attribute to return.
@@ -68,7 +69,7 @@ def __getattr__(self, attr: str) -> t.Any:
             except ConfigurationError as exception:
                 self.ctx.fail(str(exception))
 
-        return super().__getattr__(attr)  # type: ignore[no-untyped-call]
+        return super().__getattr__(attr)
 
 
 class VerdiContext(click.Context):

src/aiida/common/extendeddicts.py

@@ -8,7 +8,12 @@
 ###########################################################################
 """Various dictionary types with extended functionality."""
 
-from collections.abc import Mapping
+from __future__ import annotations
+
+from collections.abc import Hashable, KeysView, Mapping
+from typing import Any
+
+from aiida.common.typing import Self
 
 from . import exceptions
 
@@ -25,7 +30,7 @@ class AttributeDict(dict):
     used.
     """
 
-    def __init__(self, dictionary=None):
+    def __init__(self, dictionary: Mapping | None = None):
         """Recursively turn the `dict` and all its nested dictionaries into `AttributeDict` instance."""
         super().__init__()
         if dictionary is None:
@@ -37,11 +42,11 @@ def __init__(self, dictionary=None):
             else:
                 self[key] = value
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         """Representation of the object."""
         return f'{self.__class__.__name__}({dict.__repr__(self)})'
 
-    def __getattr__(self, attr):
+    def __getattr__(self, attr: Hashable) -> Any:
         """Read a key as an attribute.
 
         :raises AttributeError: if the attribute does not correspond to an existing key.
@@ -52,7 +57,7 @@ def __getattr__(self, attr):
             errmsg = f"'{self.__class__.__name__}' object has no attribute '{attr}'"
             raise AttributeError(errmsg)
 
-    def __setattr__(self, attr, value):
+    def __setattr__(self, attr: str, value: Any) -> None:
         """Set a key as an attribute."""
         try:
             self[attr] = value
@@ -61,7 +66,7 @@ def __setattr__(self, attr, value):
                 f"AttributeError: '{attr}' is not a valid attribute of the object '{self.__class__.__name__}'"
             )
 
-    def __delattr__(self, attr):
+    def __delattr__(self, attr: Hashable) -> None:
         """Delete a key as an attribute.
 
         :raises AttributeError: if the attribute does not correspond to an existing key.
@@ -72,7 +77,7 @@ def __delattr__(self, attr):
             errmsg = f"'{self.__class__.__name__}' object has no attribute '{attr}'"
             raise AttributeError(errmsg)
 
-    def __deepcopy__(self, memo=None):
+    def __deepcopy__(self, memo: Mapping | None = None) -> Self:
         """Deep copy."""
         from copy import deepcopy
 
@@ -81,15 +86,15 @@ def __deepcopy__(self, memo=None):
         retval = deepcopy(dict(self))
         return self.__class__(retval)
 
-    def __getstate__(self):
+    def __getstate__(self) -> dict[str, Any]:
         """Needed for pickling this class."""
         return self.__dict__.copy()
 
-    def __setstate__(self, dictionary):
+    def __setstate__(self, dictionary: Mapping) -> None:
         """Needed for pickling this class."""
         self.__dict__.update(dictionary)
 
-    def __dir__(self):
+    def __dir__(self) -> KeysView[Any]:
         return self.keys()
 
 
@@ -104,9 +109,9 @@ class TestExample(FixedFieldsAttributeDict):
             _valid_fields = ('a','b','c')
     """
 
-    _valid_fields = tuple()
+    _valid_fields: tuple[Any, ...] = tuple()
 
-    def __init__(self, init=None):
+    def __init__(self, init: Mapping | None = None):
         if init is None:
             init = {}
 
@@ -116,26 +121,28 @@ def __init__(self, init=None):
                 raise KeyError(errmsg)
         super().__init__(init)
 
-    def __setitem__(self, item, value):
+    def __setitem__(self, item: Hashable, value: Any) -> None:
         """Set a key as an attribute."""
         if item not in self._valid_fields:
             errmsg = f"'{item}' is not a valid key for object '{self.__class__.__name__}'"
             raise KeyError(errmsg)
         super().__setitem__(item, value)
 
-    def __setattr__(self, attr, value):
+    def __setattr__(self, attr: str, value: Any) -> None:
         """Overridden to allow direct access to fields with underscore."""
         if attr.startswith('_'):
             object.__setattr__(self, attr, value)
         else:
             super().__setattr__(attr, value)
 
     @classmethod
-    def get_valid_fields(cls):
+    def get_valid_fields(cls) -> tuple:
         """Return the list of valid fields."""
         return cls._valid_fields
 
-    def __dir__(self):
+    # TODO: We're in violation of the `dict` interface here,
+    # we should be returning collections.abc.KeysView[Any]
+    def __dir__(self) -> list[Any]:  # type: ignore[override]
         return list(self._valid_fields)
 
 
@@ -192,9 +199,9 @@ class TestExample(DefaultFieldsAttributeDict):
         See if we want that setting a default field to None means deleting it.
     """
 
-    _default_fields = tuple()
+    _default_fields: tuple = tuple()
 
-    def validate(self):
+    def validate(self) -> None:
         """Validate the keys, if any ``validate_*`` method is available."""
         for key in self.get_default_fields():
             # I get the attribute starting with validate_ and containing the name of the key
@@ -206,14 +213,14 @@ def validate(self):
                 except Exception as exc:
                     raise exceptions.ValidationError(f"Invalid value for key '{key}' [{exc.__class__.__name__}]: {exc}")
 
-    def __setattr__(self, attr, value):
+    def __setattr__(self, attr: str, value: Any) -> None:
         """Overridden to allow direct access to fields with underscore."""
         if attr.startswith('_'):
             object.__setattr__(self, attr, value)
         else:
             super().__setattr__(attr, value)
 
-    def __getitem__(self, key):
+    def __getitem__(self, key: Hashable) -> Any | None:
         """Return None instead of raising an exception if the key does not exist
         but is in the list of default fields.
         """
@@ -225,14 +232,14 @@ def __getitem__(self, key):
             raise
 
     @classmethod
-    def get_default_fields(cls):
+    def get_default_fields(cls) -> list:
         """Return the list of default fields, either defined in the instance or not."""
         return list(cls._default_fields)
 
-    def defaultkeys(self):
+    def defaultkeys(self) -> list:
         """Return the default keys defined in the instance."""
         return [_ for _ in self.keys() if _ in self._default_fields]
 
-    def extrakeys(self):
+    def extrakeys(self) -> list:
         """Return the extra keys defined in the instance."""
         return [_ for _ in self.keys() if _ not in self._default_fields]

src/aiida/common/pydantic.py

@@ -14,7 +14,7 @@
     from aiida.orm import Entity
 
 
-def get_metadata(field_info, key: str, default: t.Any | None = None):
+def get_metadata(field_info: t.Any, key: str, default: t.Any | None = None) -> t.Any:
     """Return a the metadata of the given field for a particular key.
 
     :param field_info: The field from which to retrieve the metadata.
@@ -34,15 +34,15 @@ def MetadataField(  # noqa: N802
     priority: int = 0,
     short_name: str | None = None,
     option_cls: t.Any | None = None,
-    orm_class: type['Entity'] | str | None = None,
-    orm_to_model: t.Callable[['Entity', Path], t.Any] | None = None,
+    orm_class: type[Entity] | str | None = None,
+    orm_to_model: t.Callable[[Entity, Path], t.Any] | None = None,
     model_to_orm: t.Callable[['BaseModel'], t.Any] | None = None,
     exclude_to_orm: bool = False,
     exclude_from_cli: bool = False,
     is_attribute: bool = True,
     is_subscriptable: bool = False,
-    **kwargs,
-):
+    **kwargs: t.Any,
+) -> t.Any:
     """Return a :class:`pydantic.fields.Field` instance with additional metadata.
 
     .. code-block:: python

src/aiida/common/typing.py

@@ -13,7 +13,12 @@
 import pathlib
 from typing import Union
 
-__all__ = ('FilePath',)
+try:
+    from typing import Self
+except ImportError:
+    from typing_extensions import Self
+
+__all__ = ('FilePath', 'Self')
 
 
 FilePath = Union[str, pathlib.PurePath]