👌 IMPROVE: Add typing to all base Entity classes

.pre-commit-config.yaml

@@ -82,6 +82,12 @@ repos:
                 aiida/orm/implementation/querybuilder.py|
                 aiida/orm/implementation/sqlalchemy/querybuilder/.*py|
                 aiida/orm/entities.py|
+                aiida/orm/authinfos.py|
+                aiida/orm/comments.py|
+                aiida/orm/computers.py|
+                aiida/orm/groups.py|
+                aiida/orm/logs.py|
+                aiida/orm/users.py|
                 aiida/orm/nodes/data/jsonable.py|
                 aiida/orm/nodes/node.py|
                 aiida/orm/nodes/process/.*py|

aiida/orm/authinfos.py

@@ -8,7 +8,7 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """Module for the `AuthInfo` ORM class."""
-from typing import Type
+from typing import TYPE_CHECKING, Any, Dict, Optional, Type
 
 from aiida.common import exceptions
 from aiida.common.lang import classproperty
@@ -17,6 +17,11 @@
 
 from . import entities, users
 
+if TYPE_CHECKING:
+    from aiida.orm import Computer, User
+    from aiida.orm.implementation import Backend, BackendAuthInfo
+    from aiida.transports import Transport
+
 __all__ = ('AuthInfo',)
 
 
@@ -35,7 +40,7 @@ def delete(self, pk: int) -> None:
         self._backend.authinfos.delete(pk)
 
 
-class AuthInfo(entities.Entity):
+class AuthInfo(entities.Entity['BackendAuthInfo']):
     """ORM class that models the authorization information that allows a `User` to connect to a `Computer`."""
 
     Collection = AuthInfoCollection
@@ -46,108 +51,92 @@ def objects(cls) -> AuthInfoCollection:  # pylint: disable=no-self-argument
 
     PROPERTY_WORKDIR = 'workdir'
 
-    def __init__(self, computer, user, backend=None) -> None:
+    def __init__(self, computer: 'Computer', user: 'User', backend: Optional['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`
+        :param backend: the backend to use for the instance, or use the default backend if None
         """
         backend = backend or get_manager().get_backend()
         model = backend.authinfos.create(computer=computer.backend_entity, user=user.backend_entity)
         super().__init__(model)
 
-    def __str__(self):
+    def __str__(self) -> str:
         if self.enabled:
             return f'AuthInfo for {self.user.email} on {self.computer.label}'
 
         return f'AuthInfo for {self.user.email} on {self.computer.label} [DISABLED]'
 
     @property
-    def enabled(self):
+    def enabled(self) -> bool:
         """Return whether this instance is enabled.
 
         :return: True if enabled, False otherwise
-        :rtype: bool
         """
         return self._backend_entity.enabled
 
     @enabled.setter
-    def enabled(self, enabled):
+    def enabled(self, enabled: bool) -> None:
         """Set the enabled state
 
         :param enabled: boolean, True to enable the instance, False to disable it
         """
         self._backend_entity.enabled = enabled
 
     @property
-    def computer(self):
-        """Return the computer associated with this instance.
-
-        :rtype: :class:`aiida.orm.computers.Computer`
-        """
+    def computer(self) -> 'Computer':
+        """Return the computer associated with this instance."""
         from . import computers  # pylint: disable=cyclic-import
         return computers.Computer.from_backend_entity(self._backend_entity.computer)
 
     @property
-    def user(self):
-        """Return the user associated with this instance.
-
-        :rtype: :class:`aiida.orm.users.User`
-        """
+    def user(self) -> 'User':
+        """Return the user associated with this instance."""
         return users.User.from_backend_entity(self._backend_entity.user)
 
-    def get_auth_params(self):
+    def get_auth_params(self) -> Dict[str, Any]:
         """Return the dictionary of authentication parameters
 
         :return: a dictionary with authentication parameters
-        :rtype: dict
         """
         return self._backend_entity.get_auth_params()
 
-    def set_auth_params(self, auth_params):
+    def set_auth_params(self, auth_params: Dict[str, Any]) -> None:
         """Set the dictionary of authentication parameters
 
         :param auth_params: a dictionary with authentication parameters
         """
         self._backend_entity.set_auth_params(auth_params)
 
-    def get_metadata(self):
+    def get_metadata(self) -> Dict[str, Any]:
         """Return the dictionary of metadata
 
         :return: a dictionary with metadata
-        :rtype: dict
         """
         return self._backend_entity.get_metadata()
 
-    def set_metadata(self, metadata):
+    def set_metadata(self, metadata: Dict[str, Any]) -> None:
         """Set the dictionary of metadata
 
         :param metadata: a dictionary with metadata
-        :type metadata: dict
         """
         self._backend_entity.set_metadata(metadata)
 
-    def get_workdir(self):
+    def get_workdir(self) -> str:
         """Return the working directory.
 
         If no explicit work directory is set for this instance, the working directory of the computer will be returned.
 
         :return: the working directory
-        :rtype: str
         """
         try:
             return self.get_metadata()[self.PROPERTY_WORKDIR]
         except KeyError:
             return self.computer.get_workdir()
 
-    def get_transport(self):
-        """Return a fully configured transport that can be used to connect to the computer set for this instance.
-
-        :rtype: :class:`aiida.transports.Transport`
-        """
+    def get_transport(self) -> 'Transport':
+        """Return a fully configured transport that can be used to connect to the computer set for this instance."""
         computer = self.computer
         transport_type = computer.transport_type
 

aiida/orm/comments.py

@@ -8,13 +8,18 @@
 # For further information please visit http://www.aiida.net               #
 ###########################################################################
 """Comment objects and functions"""
-from typing import List, Type
+from datetime import datetime
+from typing import TYPE_CHECKING, List, Optional, Type
 
 from aiida.common.lang import classproperty
 from aiida.manage.manager import get_manager
 
 from . import entities, users
 
+if TYPE_CHECKING:
+    from aiida.orm import Node, User
+    from aiida.orm.implementation import Backend, BackendComment
+
 __all__ = ('Comment',)
 
 
@@ -59,7 +64,7 @@ def delete_many(self, filters) -> List[int]:
         return self._backend.comments.delete_many(filters)
 
 
-class Comment(entities.Entity):
+class Comment(entities.Entity['BackendComment']):
     """Base class to map a DbComment that represents a comment attached to a certain Node."""
 
     Collection = CommentCollection
@@ -68,55 +73,59 @@ class Comment(entities.Entity):
     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):
-        """
-        Create a Comment for a given node and user
+    def __init__(self, node: 'Node', user: 'User', content: Optional[str] = None, backend: Optional['Backend'] = None):
+        """Create a Comment for a given node and user
 
         :param node: a Node instance
-        :type node: :class:`aiida.orm.Node`
-
         :param user: a User instance
-        :type user: :class:`aiida.orm.User`
-
         :param content: the comment content
-        :type content: str
+        :param backend: the backend to use for the instance, or use the default backend if None
 
         :return: a Comment object associated to the given node and user
-        :rtype: :class:`aiida.orm.Comment`
         """
         backend = backend or get_manager().get_backend()
         model = backend.comments.create(node=node.backend_entity, user=user.backend_entity, content=content)
         super().__init__(model)
 
-    def __str__(self):
+    def __str__(self) -> str:
         arguments = [self.uuid, self.node.pk, self.user.email, self.content]
         return 'Comment<{}> for node<{}> and user<{}>: {}'.format(*arguments)
 
     @property
-    def ctime(self):
+    def uuid(self) -> str:
+        """Return the UUID for this comment.
+
+        This identifier is unique across all entities types and backend instances.
+
+        :return: the entity uuid
+        """
+        return self._backend_entity.uuid
+
+    @property
+    def ctime(self) -> datetime:
         return self._backend_entity.ctime
 
     @property
-    def mtime(self):
+    def mtime(self) -> datetime:
         return self._backend_entity.mtime
 
-    def set_mtime(self, value):
+    def set_mtime(self, value: datetime) -> None:
         return self._backend_entity.set_mtime(value)
 
     @property
-    def node(self):
+    def node(self) -> 'Node':
         return self._backend_entity.node
 
     @property
-    def user(self):
+    def user(self) -> 'User':
         return users.User.from_backend_entity(self._backend_entity.user)
 
-    def set_user(self, value):
+    def set_user(self, value: 'User') -> None:
         self._backend_entity.user = value.backend_entity
 
     @property
-    def content(self):
+    def content(self) -> str:
         return self._backend_entity.content
 
-    def set_content(self, value):
+    def set_content(self, value: str) -> None:
         return self._backend_entity.set_content(value)