👌 IMPROVE: Entity Collection typing (#5183)

This PR makes static types available for `Entity.objects -> Collection`.
The `Collection` subclasses are now defined outside of the `Entity` subclasses,
using `Generic`s to dynamically define the entity subclass their methods return.
A number of fixes were also made, to ensure the correct backend is used in queries.

.pre-commit-config.yaml

@@ -81,6 +81,7 @@ repos:
                 aiida/orm/implementation/sqlalchemy/backend.py|
                 aiida/orm/implementation/querybuilder.py|
                 aiida/orm/implementation/sqlalchemy/querybuilder/.*py|
+                aiida/orm/entities.py|
                 aiida/orm/nodes/data/jsonable.py|
                 aiida/orm/nodes/node.py|
                 aiida/orm/nodes/process/.*py|

aiida/common/datastructures.py

@@ -165,39 +165,3 @@ class CodeRunMode(IntEnum):
 
     SERIAL = 0
     PARALLEL = 1
-
-
-class LazyStore:
-    """
-    A container that provides a mapping to objects based on a key, if the object is not
-    found in the container when it is retrieved it will created using a provided factory
-    method
-    """
-
-    def __init__(self):
-        self._store = {}
-
-    def get(self, key, factory):
-        """
-        Get a value in the store based on the key, if it doesn't exist it will be created
-        using the factory method and returned
-
-        :param key: the key of the object to get
-        :param factory: the factory used to create the object if necessary
-        :return: the object
-        """
-        try:
-            return self._store[key]
-        except KeyError:
-            obj = factory()
-            self._store[key] = obj
-            return obj
-
-    def pop(self, key):
-        """
-        Pop an object from the store based on the given key
-
-        :param key: the object key
-        :return: the object that was popped
-        """
-        return self._store.pop(key)

aiida/common/lang.py

@@ -11,7 +11,7 @@
 import functools
 import inspect
 import keyword
-from typing import Any, Callable, Generic, TypeVar
+from typing import Any, Callable, Generic, Type, TypeVar
 
 
 def isidentifier(identifier):
@@ -77,6 +77,7 @@ def wrapped_fn(self, *args, **kwargs):  # pylint: disable=missing-docstring
 override = override_decorator(check=False)  # pylint: disable=invalid-name
 
 ReturnType = TypeVar('ReturnType')
+SelfType = TypeVar('SelfType')
 
 
 class classproperty(Generic[ReturnType]):  # pylint: disable=invalid-name
@@ -88,8 +89,8 @@ class classproperty(Generic[ReturnType]):  # pylint: disable=invalid-name
     instance as its first argument).
     """
 
-    def __init__(self, getter: Callable[[Any], ReturnType]) -> None:
+    def __init__(self, getter: Callable[[Type[SelfType]], ReturnType]) -> None:
         self.getter = getter
 
-    def __get__(self, instance: Any, owner: Any) -> ReturnType:
+    def __get__(self, instance: Any, owner: Type[SelfType]) -> ReturnType:
         return self.getter(owner)

aiida/orm/authinfos.py

@@ -8,8 +8,10 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """Module for the `AuthInfo` ORM class."""
+from typing import Type
 
 from aiida.common import exceptions
+from aiida.common.lang import classproperty
 from aiida.manage.manager import get_manager
 from aiida.plugins import TransportFactory
 
@@ -18,31 +20,40 @@
 __all__ = ('AuthInfo',)
 
 
+class AuthInfoCollection(entities.Collection['AuthInfo']):
+    """The collection of `AuthInfo` entries."""
+
+    @staticmethod
+    def _entity_base_cls() -> Type['AuthInfo']:
+        return AuthInfo
+
+    def delete(self, pk: int) -> None:
+        """Delete an entry from the collection.
+
+        :param pk: the pk of the entry to delete
+        """
+        self._backend.authinfos.delete(pk)
+
+
 class AuthInfo(entities.Entity):
     """ORM class that models the authorization information that allows a `User` to connect to a `Computer`."""
 
-    class Collection(entities.Collection):
-        """The collection of `AuthInfo` entries."""
+    Collection = AuthInfoCollection
 
-        def delete(self, pk):
-            """Delete an entry from the collection.
-
-            :param pk: the pk of the entry to delete
-            """
-            self._backend.authinfos.delete(pk)
+    @classproperty
+    def objects(cls) -> AuthInfoCollection:  # pylint: disable=no-self-argument
+        return AuthInfoCollection.get_cached(cls, get_manager().get_backend())
 
     PROPERTY_WORKDIR = 'workdir'
 
-    def __init__(self, computer, user, backend=None):
+    def __init__(self, computer, user, backend=None) -> None:
         """Create an `AuthInfo` instance for the given computer and user.
 
         :param computer: a `Computer` instance
         :type computer: :class:`aiida.orm.Computer`
 
         :param user: a `User` instance
         :type user: :class:`aiida.orm.User`
-
-        :rtype: :class:`aiida.orm.authinfos.AuthInfo`
         """
         backend = backend or get_manager().get_backend()
         model = backend.authinfos.create(computer=computer.backend_entity, user=user.backend_entity)

aiida/orm/comments.py

@@ -8,54 +8,65 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """Comment objects and functions"""
+from typing import List, Type
 
+from aiida.common.lang import classproperty
 from aiida.manage.manager import get_manager
 
 from . import entities, users
 
 __all__ = ('Comment',)
 
 
-class Comment(entities.Entity):
-    """Base class to map a DbComment that represents a comment attached to a certain Node."""
+class CommentCollection(entities.Collection['Comment']):
+    """The collection of Comment entries."""
+
+    @staticmethod
+    def _entity_base_cls() -> Type['Comment']:
+        return Comment
+
+    def delete(self, pk: int) -> None:
+        """
+        Remove a Comment from the collection with the given id
+
+        :param pk: the id of the comment to delete
+
+        :raises TypeError: if ``comment_id`` is not an `int`
+        :raises `~aiida.common.exceptions.NotExistent`: if Comment with ID ``comment_id`` is not found
+        """
+        self._backend.comments.delete(pk)
 
-    class Collection(entities.Collection):
-        """The collection of Comment entries."""
+    def delete_all(self) -> None:
+        """
+        Delete all Comments from the Collection
 
-        def delete(self, comment_id):
-            """
-            Remove a Comment from the collection with the given id
+        :raises `~aiida.common.exceptions.IntegrityError`: if all Comments could not be deleted
+        """
+        self._backend.comments.delete_all()
 
-            :param comment_id: the id of the comment to delete
-            :type comment_id: int
+    def delete_many(self, filters) -> List[int]:
+        """
+        Delete Comments from the Collection based on ``filters``
 
-            :raises TypeError: if ``comment_id`` is not an `int`
-            :raises `~aiida.common.exceptions.NotExistent`: if Comment with ID ``comment_id`` is not found
-            """
-            self._backend.comments.delete(comment_id)
+        :param filters: similar to QueryBuilder filter
+        :type filters: dict
 
-        def delete_all(self):
-            """
-            Delete all Comments from the Collection
+        :return: (former) ``PK`` s of deleted Comments
 
-            :raises `~aiida.common.exceptions.IntegrityError`: if all Comments could not be deleted
-            """
-            self._backend.comments.delete_all()
+        :raises TypeError: if ``filters`` is not a `dict`
+        :raises `~aiida.common.exceptions.ValidationError`: if ``filters`` is empty
+        """
+        return self._backend.comments.delete_many(filters)
 
-        def delete_many(self, filters):
-            """
-            Delete Comments from the Collection based on ``filters``
 
-            :param filters: similar to QueryBuilder filter
-            :type filters: dict
+class Comment(entities.Entity):
+    """Base class to map a DbComment that represents a comment attached to a certain Node."""
 
-            :return: (former) ``PK`` s of deleted Comments
-            :rtype: list
+    Collection = CommentCollection
 
-            :raises TypeError: if ``filters`` is not a `dict`
-            :raises `~aiida.common.exceptions.ValidationError`: if ``filters`` is empty
-            """
-            self._backend.comments.delete_many(filters)
+    @classproperty
+    def objects(cls) -> CommentCollection:  # pylint: disable=no-self-argument
+        return CommentCollection.get_cached(cls, get_manager().get_backend())
 
     def __init__(self, node, user, content=None, backend=None):
         """

aiida/orm/computers.py

@@ -10,8 +10,10 @@
 """Module for Computer entities"""
 import logging
 import os
+from typing import List, Optional, Tuple, Type
 
 from aiida.common import exceptions
+from aiida.common.lang import classproperty
 from aiida.manage.manager import get_manager
 from aiida.orm.implementation import Backend
 from aiida.plugins import SchedulerFactory, TransportFactory
@@ -21,6 +23,40 @@
 __all__ = ('Computer',)
 
 
+class ComputerCollection(entities.Collection['Computer']):
+    """The collection of Computer entries."""
+
+    @staticmethod
+    def _entity_base_cls() -> Type['Computer']:
+        return Computer
+
+    def get_or_create(self, label: Optional[str] = None, **kwargs) -> Tuple['Computer', bool]:
+        """
+        Try to retrieve a Computer from the DB with the given arguments;
+        create (and store) a new Computer if such a Computer was not present yet.
+
+        :param label: computer label
+
+        :return: (computer, created) where computer is the computer (new or existing,
+            in any case already stored) and created is a boolean saying
+        """
+        if not label:
+            raise ValueError('Computer label must be provided')
+
+        try:
+            return False, self.get(label=label)
+        except exceptions.NotExistent:
+            return True, Computer(backend=self.backend, label=label, **kwargs)
+
+    def list_labels(self) -> List[str]:
+        """Return a list with all the labels of the computers in the DB."""
+        return self._backend.computers.list_names()
+
+    def delete(self, pk: int) -> None:
+        """Delete the computer with the given id"""
+        return self._backend.computers.delete(pk)
+
+
 class Computer(entities.Entity):
     """
     Computer entity.
@@ -34,36 +70,11 @@ class Computer(entities.Entity):
     PROPERTY_WORKDIR = 'workdir'
     PROPERTY_SHEBANG = 'shebang'
 
-    class Collection(entities.Collection):
-        """The collection of Computer entries."""
-
-        def get_or_create(self, label=None, **kwargs):
-            """
-            Try to retrieve a Computer from the DB with the given arguments;
-            create (and store) a new Computer if such a Computer was not present yet.
-
-            :param label: computer label
-            :type label: str
-
-            :return: (computer, created) where computer is the computer (new or existing,
-              in any case already stored) and created is a boolean saying
-            :rtype: (:class:`aiida.orm.Computer`, bool)
-            """
-            if not label:
-                raise ValueError('Computer label must be provided')
-
-            try:
-                return False, self.get(label=label)
-            except exceptions.NotExistent:
-                return True, Computer(backend=self.backend, label=label, **kwargs)
-
-        def list_labels(self):
-            """Return a list with all the labels of the computers in the DB."""
-            return self._backend.computers.list_names()
+    Collection = ComputerCollection
 
-        def delete(self, id):  # pylint: disable=redefined-builtin,invalid-name
-            """Delete the computer with the given id"""
-            return self._backend.computers.delete(id)
+    @classproperty
+    def objects(cls) -> ComputerCollection:  # pylint: disable=no-self-argument
+        return ComputerCollection.get_cached(cls, get_manager().get_backend())
 
     def __init__(  # pylint: disable=too-many-arguments
         self,