Refactor Model classes out of Entity classes

Author: edan-bainglass

docs/source/nitpick-exceptions

@@ -33,8 +33,9 @@ py:meth tempfile.TemporaryDirectory
 ### AiiDA
 py:class ReturnType
 py:class SelfType
-py:class CollectionType
+py:class EntityCollectionType
 py:class EntityType
+py:class EntityModelType
 py:class EntityClsType
 py:class ProjectType
 py:class BackendEntityType
@@ -51,8 +52,9 @@ py:class UserDefinedType
 py:class LinkAnnotateType
 py:obj ReturnType
 py:obj SelfType
-py:obj CollectionType
+py:obj EntityCollectionType
 py:obj EntityType
+py:obj EntityModelType
 py:obj BackendEntityType
 py:obj BackendNodeType
 py:obj TransactionType
@@ -62,23 +64,19 @@ py:class aiida.common.ReturnType
 py:class aiida.common.SelfType
 py:class aiida.common.lang.MethodType
 py:class aiida.cmdline.params._shims.C
-py:class aiida.orm.entitites.CollectionType
-py:class aiida.orm.entitites.EntityType
-py:class aiida.orm.entitites.BackendEntityType
 py:class aiida.orm.groups.SelfType
-py:class aiida.orm.implementation.entitites.EntityType
 py:class aiida.engine.processes.functions.FunctionType
 py:class aiida.engine.processes.workchains.workchain.MethodType
 py:class aiida.orm.entities.EntityInputModel
-py:class aiida.orm.entities.EntityModelType
 py:class aiida.orm.entities.EntityType
+py:class aiida.orm.entities.EntityModelType
 py:class aiida.orm.entities.BackendEntityType
-py:class aiida.orm.entities.CollectionType
+py:class aiida.orm.entities.EntityCollectionType
 py:class aiida.orm.implementation.nodes.BackendNodeType
 py:class aiida.orm.implementation.storage_backend.TransactionType
 py:class aiida.orm.implementation.entities.EntityType
 py:class aiida.orm.nodes.data.enum.EnumType
-py:class aiida.orm.nodes.NodeType
+py:class aiida.orm.nodes.node.NodeType
 py:class aiida.storage.psql_dos.orm.ModelType
 py:class aiida.storage.psql_dos.orm.SelfType
 py:class aiida.storage.psql_dos.orm.entities.SelfType
@@ -91,18 +89,15 @@ py:class aiida.common.lang.ReturnType
 py:obj aiida.common.ReturnType
 py:obj aiida.common.SelfType
 py:obj aiida.common.lang.ReturnType
-py:obj aiida.orm.entitites.CollectionType
-py:obj aiida.orm.entitites.EntityType
-py:obj aiida.orm.entitites.BackendEntityType
 py:obj aiida.orm.groups.SelfType
-py:obj aiida.orm.entities.CollectionType
+py:obj aiida.orm.entities.EntityCollectionType
 py:obj aiida.orm.entities.BackendEntityType
 py:obj aiida.orm.entities.EntityType
+py:obj aiida.orm.entities.EntityModelType
 py:obj aiida.orm.implementation.entities.EntityType
 py:obj aiida.orm.implementation.nodes.BackendNodeType
 py:obj aiida.orm.implementation.storage_backend.TransactionType
 py:obj aiida.orm.nodes.node.NodeType
-py:obj aiida.orm.nodes.NodeType
 py:obj aiida.storage.psql_dos.orm.ModelType
 py:obj aiida.storage.psql_dos.orm.SelfType
 py:obj aiida.storage.psql_dos.orm.entities.ModelType

pyproject.toml

@@ -479,7 +479,7 @@ select = [
 # Mark some classes as generic, per https://docs.astral.sh/ruff/settings/#lint_pyflakes_extend-generics
 # Needed due to https://github.com/astral-sh/ruff/issues/9298
 [tool.ruff.lint.pyflakes]
-extend-generics = ["aiida.orm.entities.Entity", "aiida.orm.entities.Collection"]
+extend-generics = ["aiida.orm.entities.Entity", "aiida.orm.entities.EntityCollection", "aiida.orm.entities.EntityModel"]
 
 [tool.tox]
 legacy_tox_ini = """

src/aiida/common/typing.py

@@ -19,6 +19,11 @@
 else:
     from typing_extensions import Self
 
-__all__ = ('FilePath', 'Self')
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+__all__ = ('FilePath', 'Self', 'TypeAlias')
 
 FilePath = Union[str, pathlib.PurePath]

src/aiida/orm/__init__.py

@@ -45,14 +45,14 @@
     'CifData',
     'Code',
     'CodeEntityLoader',
-    'Collection',
     'Comment',
     'Computer',
     'ComputerEntityLoader',
     'ContainerizedCode',
     'Data',
     'Dict',
     'Entity',
+    'EntityCollection',
     'EntityExtras',
     'EntityTypes',
     'EnumData',

src/aiida/orm/authinfos.py

@@ -14,6 +14,7 @@
 
 from aiida.common import exceptions
 from aiida.common.pydantic import MetadataField
+from aiida.common.typing import TypeAlias
 from aiida.manage import get_manager
 from aiida.plugins import TransportFactory
 
@@ -29,7 +30,7 @@
 __all__ = ('AuthInfo',)
 
 
-class AuthInfoCollection(entities.Collection['AuthInfo']):
+class AuthInfoCollection(entities.EntityCollection['AuthInfo']):
     """The collection of `AuthInfo` entries."""
 
     @staticmethod
@@ -44,40 +45,44 @@ def delete(self, pk: int) -> None:
         self._backend.authinfos.delete(pk)
 
 
+class AuthInfoModel(entities.Entity.Model):
+    computer: int = MetadataField(
+        description='The PK of the computer',
+        is_attribute=False,
+        orm_class=Computer,
+        orm_to_model=lambda auth_info: cast(AuthInfo, auth_info).computer.pk,
+    )
+    user: int = MetadataField(
+        description='The PK of the user',
+        is_attribute=False,
+        orm_class=User,
+        orm_to_model=lambda auth_info: cast(AuthInfo, auth_info).user.pk,
+    )
+    enabled: bool = MetadataField(
+        True,
+        description='Whether the instance is enabled',
+        is_attribute=False,
+    )
+    auth_params: Dict[str, Any] = MetadataField(
+        default_factory=dict,
+        description='Dictionary of authentication parameters',
+        is_attribute=False,
+    )
+    metadata: Dict[str, Any] = MetadataField(
+        default_factory=dict,
+        description='Dictionary of metadata',
+        is_attribute=False,
+    )
+
+
 class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection]):
     """ORM class that models the authorization information that allows a `User` to connect to a `Computer`."""
 
+    Model: TypeAlias = AuthInfoModel
+
     _CLS_COLLECTION = AuthInfoCollection
-    PROPERTY_WORKDIR = 'workdir'
 
-    class Model(entities.Entity.Model):
-        computer: int = MetadataField(
-            description='The PK of the computer',
-            is_attribute=False,
-            orm_class=Computer,
-            orm_to_model=lambda auth_info: cast(AuthInfo, auth_info).computer.pk,
-        )
-        user: int = MetadataField(
-            description='The PK of the user',
-            is_attribute=False,
-            orm_class=User,
-            orm_to_model=lambda auth_info: cast(AuthInfo, auth_info).user.pk,
-        )
-        enabled: bool = MetadataField(
-            True,
-            description='Whether the instance is enabled',
-            is_attribute=False,
-        )
-        auth_params: Dict[str, Any] = MetadataField(
-            default_factory=dict,
-            description='Dictionary of authentication parameters',
-            is_attribute=False,
-        )
-        metadata: Dict[str, Any] = MetadataField(
-            default_factory=dict,
-            description='Dictionary of metadata',
-            is_attribute=False,
-        )
+    PROPERTY_WORKDIR = 'workdir'
 
     def __init__(
         self,

src/aiida/orm/comments.py

@@ -15,6 +15,7 @@
 from uuid import UUID
 
 from aiida.common.pydantic import MetadataField
+from aiida.common.typing import TypeAlias
 from aiida.manage import get_manager
 
 from . import entities
@@ -28,7 +29,7 @@
 __all__ = ('Comment',)
 
 
-class CommentCollection(entities.Collection['Comment']):
+class CommentCollection(entities.EntityCollection['Comment']):
     """The collection of Comment entries."""
 
     @staticmethod
@@ -65,43 +66,46 @@ def delete_many(self, filters: dict) -> List[int]:
         return self._backend.comments.delete_many(filters)
 
 
+class CommentModel(entities.Entity.Model):
+    uuid: UUID = MetadataField(
+        description='The UUID of the comment',
+        is_attribute=False,
+        exclude_to_orm=True,
+    )
+    ctime: datetime = MetadataField(
+        description='Creation time of the comment',
+        is_attribute=False,
+        exclude_to_orm=True,
+    )
+    mtime: datetime = MetadataField(
+        description='Modified time of the comment',
+        is_attribute=False,
+        exclude_to_orm=True,
+    )
+    node: int = MetadataField(
+        description='Node PK that the comment is attached to',
+        is_attribute=False,
+        orm_class='core.node',
+        orm_to_model=lambda comment: cast(Comment, comment).node.pk,
+    )
+    user: int = MetadataField(
+        description='User PK that created the comment',
+        is_attribute=False,
+        orm_class='core.user',
+        orm_to_model=lambda comment: cast(Comment, comment).user.pk,
+    )
+    content: str = MetadataField(
+        description='Content of the comment',
+        is_attribute=False,
+    )
+
+
 class Comment(entities.Entity['BackendComment', CommentCollection]):
     """Base class to map a DbComment that represents a comment attached to a certain Node."""
 
-    _CLS_COLLECTION = CommentCollection
+    Model: TypeAlias = CommentModel
 
-    class Model(entities.Entity.Model):
-        uuid: UUID = MetadataField(
-            description='The UUID of the comment',
-            is_attribute=False,
-            exclude_to_orm=True,
-        )
-        ctime: datetime = MetadataField(
-            description='Creation time of the comment',
-            is_attribute=False,
-            exclude_to_orm=True,
-        )
-        mtime: datetime = MetadataField(
-            description='Modified time of the comment',
-            is_attribute=False,
-            exclude_to_orm=True,
-        )
-        node: int = MetadataField(
-            description='Node PK that the comment is attached to',
-            is_attribute=False,
-            orm_class='core.node',
-            orm_to_model=lambda comment: cast(Comment, comment).node.pk,
-        )
-        user: int = MetadataField(
-            description='User PK that created the comment',
-            is_attribute=False,
-            orm_class='core.user',
-            orm_to_model=lambda comment: cast(Comment, comment).user.pk,
-        )
-        content: str = MetadataField(
-            description='Content of the comment',
-            is_attribute=False,
-        )
+    _CLS_COLLECTION = CommentCollection
 
     def __init__(
         self, node: 'Node', user: 'User', content: Optional[str] = None, backend: Optional['StorageBackend'] = None

src/aiida/orm/computers.py

@@ -17,6 +17,7 @@
 from aiida.common import exceptions
 from aiida.common.log import AIIDA_LOGGER, AiidaLoggerType
 from aiida.common.pydantic import MetadataField
+from aiida.common.typing import TypeAlias
 from aiida.manage import get_manager
 from aiida.plugins import SchedulerFactory, TransportFactory
 
@@ -31,7 +32,7 @@
 __all__ = ('Computer',)
 
 
-class ComputerCollection(entities.Collection['Computer']):
+class ComputerCollection(entities.EntityCollection['Computer']):
     """The collection of Computer entries."""
 
     @staticmethod
@@ -64,9 +65,45 @@ def delete(self, pk: int) -> None:
         return self._backend.computers.delete(pk)
 
 
+class ComputerModel(entities.Entity.Model):
+    uuid: UUID = MetadataField(
+        description='The UUID of the computer',
+        is_attribute=False,
+        exclude_to_orm=True,
+    )
+    label: str = MetadataField(
+        description='Label for the computer',
+        is_attribute=False,
+    )
+    description: str = MetadataField(
+        '',
+        description='Description of the computer',
+        is_attribute=False,
+    )
+    hostname: str = MetadataField(
+        description='Hostname of the computer',
+        is_attribute=False,
+    )
+    transport_type: str = MetadataField(
+        description='Transport type of the computer',
+        is_attribute=False,
+    )
+    scheduler_type: str = MetadataField(
+        description='Scheduler type of the computer',
+        is_attribute=False,
+    )
+    metadata: Dict[str, Any] = MetadataField(
+        default_factory=dict,
+        description='Metadata of the computer',
+        is_attribute=False,
+    )
+
+
 class Computer(entities.Entity['BackendComputer', ComputerCollection]):
     """Computer entity."""
 
+    Model: TypeAlias = ComputerModel
+
     _logger = AIIDA_LOGGER.getChild('orm.computers')
 
     PROPERTY_MINIMUM_SCHEDULER_POLL_INTERVAL = 'minimum_scheduler_poll_interval'
@@ -76,39 +113,6 @@ class Computer(entities.Entity['BackendComputer', ComputerCollection]):
 
     _CLS_COLLECTION = ComputerCollection
 
-    class Model(entities.Entity.Model):
-        uuid: UUID = MetadataField(
-            description='The UUID of the computer',
-            is_attribute=False,
-            exclude_to_orm=True,
-        )
-        label: str = MetadataField(
-            description='Label for the computer',
-            is_attribute=False,
-        )
-        description: str = MetadataField(
-            '',
-            description='Description of the computer',
-            is_attribute=False,
-        )
-        hostname: str = MetadataField(
-            description='Hostname of the computer',
-            is_attribute=False,
-        )
-        transport_type: str = MetadataField(
-            description='Transport type of the computer',
-            is_attribute=False,
-        )
-        scheduler_type: str = MetadataField(
-            description='Scheduler type of the computer',
-            is_attribute=False,
-        )
-        metadata: Dict[str, Any] = MetadataField(
-            default_factory=dict,
-            description='Metadata of the computer',
-            is_attribute=False,
-        )
-
     def __init__(
         self,
         label: Optional[str] = None,