👌 IMPROVE: Entity Collection typing

[chrisjsewell]

This PR makes static types available for Entity.ojects.
As an example, in pyproject.toml comment out follow_imports:

[[tool.mypy.overrides]]
module = 'aiida.*'
# can only follow these imports when more of the code is typed
# follow_imports = "skip"

and run example.py against mypy (with the pre-commit env loaded):

from aiida.orm import Computer, Data, Dict, Group, User
from aiida.manage.manager import get_manager

reveal_type(Computer.objects)
reveal_type(Group.objects.all())
reveal_type(User.objects.get_default())
reveal_type(Data.objects)
reveal_type(Data.objects(get_manager().get_backend()))
reveal_type(Data.objects.get())
reveal_type(Dict.objects.get())

data = Data.objects.all()[0]
reveal_type(data)

(.../aiida_core_develop/.tox/py38-pre-commit) aiida_core_develop/$ mypy example.py
example.py:4: note: Revealed type is "aiida.orm.computers.ComputerCollection"
example.py:5: note: Revealed type is "builtins.list[aiida.orm.groups.Group*]"
example.py:6: note: Revealed type is "aiida.orm.users.User"
example.py:7: note: Revealed type is "aiida.orm.nodes.node.NodeCollection[aiida.orm.nodes.data.data.Data*]"
example.py:8: note: Revealed type is "aiida.orm.nodes.node.NodeCollection[aiida.orm.nodes.data.data.Data]"
example.py:9: note: Revealed type is "aiida.orm.nodes.data.data.Data*"
example.py:10: note: Revealed type is "aiida.orm.nodes.data.dict.Dict*"
example.py:13: note: Revealed type is "aiida.orm.nodes.data.data.Data*"

One thing to note above, is that entity subclass methods give base-class entity types, Data.objects.get() -> Node.
Ideally this would give Data, but I haven't been able to work this out (it is already tricky enough, with the cyclic nature of Entity -> Collection -> Entity ... 😬)
To stress, this is only an issue with the static typing, not the actual runtime return types. (fixed, see below)

A few points:

• The Collection sub-classes are now defined outside of the Entity sub-classes. There is no reason why they cannot be used independent of the Entity itself.
• I left the Entity.Collection attribute, for back-compatibility, but actually it is no longer needed.
• The aiida/common/datastructures::LazyStore has been removed and replaced with https://docs.python.org/3/library/functools.html#functools.lru_cache. Caching was added due to Add caching of object collections #2216; I'm interested if it really speeds things up that much, but also there's no downside to keeping it.
  You can view it working with:

  In [1]: Data.objects; Data.objects; User.objects; Group.objects.get_cached.cache_info()
  Out[1]: CacheInfo(hits=1, misses=3, maxsize=100, currsize=3)

• There was a number of uses of e.g. Comment.objects that were fixed to ensure the Collection uses the same backend as the Entity, e.g. Comment.objects(self.backend)
• 🔧 MAINTAIN: Add pylint_aiida plugin #5182 ensures e.g. Node.objects.get uses don't raise (incorrect) no-member issues in pylint
• I noted that User.get_or_create returns the opposite way to Computer or Group, i.e. Tuple[bool, 'User'], not Tuple['User', bool]. But I don't see away to deprecate that 🤷

[codecov]

Codecov Report

Merging #5183 (0ae56a1) into develop (ede283e) will increase coverage by 0.03%.
The diff coverage is 94.84%.

@@             Coverage Diff             @@
##           develop    #5183      +/-   ##
===========================================
+ Coverage    80.97%   81.00%   +0.03%     
===========================================
  Files          535      535              
  Lines        37201    37289      +88     
===========================================
+ Hits         30120    30201      +81     
- Misses        7081     7088       +7     

Flag	Coverage Δ
django	75.80% <94.84%> (+0.05%)	⬆️
sqlalchemy	74.91% <94.84%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Impacted Files	Coverage Δ
aiida/common/datastructures.py	100.00% <ø> (+3.13%)	⬆️
aiida/orm/computers.py	81.95% <90.48%> (+0.61%)	⬆️
aiida/orm/logs.py	93.16% <90.48%> (-1.84%)	⬇️
aiida/orm/groups.py	95.05% <91.31%> (+0.31%)	⬆️
aiida/orm/entities.py	96.69% <95.24%> (-1.64%)	⬇️
aiida/orm/nodes/node.py	96.33% <95.84%> (+0.04%)	⬆️
aiida/orm/users.py	85.89% <97.23%> (+0.75%)	⬆️
aiida/common/lang.py	80.49% <100.00%> (+0.49%)	⬆️
aiida/orm/authinfos.py	86.21% <100.00%> (+2.54%)	⬆️
aiida/orm/comments.py	93.19% <100.00%> (+1.76%)	⬆️
... and 4 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update ede283e...0ae56a1. Read the comment docs.

[chrisjsewell]

Note, you can do:

NodeType = TypeVar('NodeType', bound='Node')

class NodeCollection(EntityCollection[NodeType], Generic[NodeType]):

Then using e.g. NodeCollection[Data] properly gives Data.objects.get() -> Data.
But doing something like:

    @classproperty
    def objects(cls: Type[NodeType]) -> NodeCollection[NodeType]:

ends up with note: Revealed type is "NodeType`-1", i.e. it can't determine cls correctly (see python/mypy#2563)

---

Edit, as noted in python/mypy#2563 (comment), if you change classproperty to:

ReturnType = TypeVar('ReturnType')
ClsType = TypeVar('ClsType')


class classproperty(Generic[ReturnType]):  # pylint: disable=invalid-name
    """
    A class that, when used as a decorator, works as if the
    two decorators @property and @classmethod where applied together
    (i.e., the object works as a property, both for the Class and for any
    of its instance; and is called with the class cls rather than with the
    instance as its first argument).
    """

    def __init__(self, getter: Callable[[Type[ClsType]], ReturnType]) -> None:
        self.getter = getter

    def __get__(self, instance: Any, owner: Type[ClsType]) -> ReturnType:
        return self.getter(owner)

It does now give note: Revealed type is "aiida.orm.nodes.data.data.Data*"

Pylance introspection also works... kind of.
You can see below that it doesn't fully resolve the node type, but does give autocompletion for methods on the Node (just not methods specific to the Data class e.g. Data.source).

---

now implemented

[sphuber]

Thanks @chrisjsewell , just some minor changes, some questions for clarification and one potential problem which was already there and your PR doesn't change, so this could be left for another time but wanted to raise it here in case you have an idea already.

[sphuber]

Thanks @chrisjsewell merge away