Discard type aliasing and inject a generic EntityModelType

Author: edan-bainglass

src/aiida/common/pydantic.py

@@ -33,9 +33,9 @@ def MetadataField(  # noqa: N802
     priority: int = 0,
     short_name: str | None = None,
     option_cls: t.Any | None = None,
-    orm_class: type[Entity[t.Any, t.Any]] | str | None = None,
-    orm_to_model: t.Callable[[Entity[t.Any, t.Any]], t.Any]  # without arguments
-    | t.Callable[[Entity[t.Any, t.Any], dict[str, t.Any]], t.Any]  # with arguments
+    orm_class: type[Entity[t.Any, t.Any, t.Any]] | str | None = None,
+    orm_to_model: t.Callable[[Entity[t.Any, t.Any, t.Any]], t.Any]  # without arguments
+    | t.Callable[[Entity[t.Any, t.Any, t.Any], dict[str, t.Any]], t.Any]  # with arguments
     | None = None,
     model_to_orm: t.Callable[[BaseModel], t.Any] | None = None,
     exclude_to_orm: bool = False,

src/aiida/orm/authinfos.py

@@ -14,7 +14,6 @@
 
 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
 
@@ -75,10 +74,10 @@ class AuthInfoModel(entities.EntityModel):
     )
 
 
-class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection]):
+class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection, AuthInfoModel]):
     """ORM class that models the authorization information that allows a `User` to connect to a `Computer`."""
 
-    Model: TypeAlias = AuthInfoModel
+    Model = AuthInfoModel
 
     _CLS_COLLECTION = AuthInfoCollection
 

src/aiida/orm/comments.py

@@ -15,7 +15,6 @@
 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
@@ -100,10 +99,10 @@ class CommentModel(entities.EntityModel):
     )
 
 
-class Comment(entities.Entity['BackendComment', CommentCollection]):
+class Comment(entities.Entity['BackendComment', CommentCollection, CommentModel]):
     """Base class to map a DbComment that represents a comment attached to a certain Node."""
 
-    Model: TypeAlias = CommentModel
+    Model = CommentModel
 
     _CLS_COLLECTION = CommentCollection
 

src/aiida/orm/computers.py

@@ -17,7 +17,6 @@
 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
 
@@ -99,10 +98,10 @@ class ComputerModel(entities.EntityModel):
     )
 
 
-class Computer(entities.Entity['BackendComputer', ComputerCollection]):
+class Computer(entities.Entity['BackendComputer', ComputerCollection, ComputerModel]):
     """Computer entity."""
 
-    Model: TypeAlias = ComputerModel
+    Model = ComputerModel
 
     _logger = AIIDA_LOGGER.getChild('orm.computers')
 

src/aiida/orm/entities.py

@@ -19,12 +19,13 @@
 from plumpy.base.utils import call_with_super_check, super_check
 from pydantic import BaseModel, ConfigDict, create_model
 from pydantic.fields import FieldInfo
+from typing_extensions import overload
 
 from aiida.common import exceptions, log
 from aiida.common.exceptions import EntryPointError, InvalidOperation, NotExistent
 from aiida.common.lang import classproperty, type_check
 from aiida.common.pydantic import MetadataField, get_metadata
-from aiida.common.typing import Self, TypeAlias
+from aiida.common.typing import Self
 from aiida.common.warnings import warn_deprecation
 from aiida.manage import get_manager
 
@@ -38,7 +39,8 @@
 
 BackendEntityType = TypeVar('BackendEntityType', bound='BackendEntity')
 EntityCollectionType = TypeVar('EntityCollectionType', bound='EntityCollection[Any]')
-EntityType = TypeVar('EntityType', bound='Entity[Any,Any]')
+EntityModelType = TypeVar('EntityModelType', bound='EntityModel')
+EntityType = TypeVar('EntityType', bound='Entity[Any,Any,Any]')
 
 
 class EntityTypes(Enum):
@@ -189,7 +191,7 @@ class EntityModel(BaseModel, defer_build=True):
     )
 
     @classmethod
-    def as_create_model(cls: Type[EntityModel]) -> Type[EntityModel]:
+    def as_create_model(cls: type[EntityModel]) -> type[BaseModel]:
         """Return a derived creation model class with read-only fields removed.
 
         This also removes any serializers/validators defined on those fields.
@@ -238,16 +240,16 @@ def _prune_field_decorators(field_decorators: dict[str, Any]) -> dict[str, Any]:
         return CreateModel
 
 
-class Entity(abc.ABC, Generic[BackendEntityType, EntityCollectionType], metaclass=EntityFieldMeta):
+class Entity(abc.ABC, Generic[BackendEntityType, EntityCollectionType, EntityModelType], metaclass=EntityFieldMeta):
     """An AiiDA entity"""
 
-    Model: TypeAlias = EntityModel
+    Model: type[EntityModelType] = EntityModel  # type: ignore[assignment]
 
     _CLS_COLLECTION: type[EntityCollectionType] = EntityCollection  # type: ignore[assignment]
     _logger = log.AIIDA_LOGGER.getChild('orm.entities')
 
     @classproperty
-    def CreateModel(cls) -> Type[Model]:  # noqa: N802, N805
+    def CreateModel(cls) -> type[BaseModel]:  # noqa: N802, N805
         """Return the creation version of the model class for this entity.
 
         :return: The creation model class, with read-only fields removed.
@@ -261,7 +263,7 @@ def model_to_orm_fields(cls) -> dict[str, FieldInfo]:
         }
 
     @classmethod
-    def model_to_orm_field_values(cls, model: EntityModel) -> dict[str, Any]:
+    def model_to_orm_field_values(cls, model: BaseModel) -> dict[str, Any]:
         from aiida.plugins.factories import BaseFactory
 
         fields = {}
@@ -291,13 +293,40 @@ def model_to_orm_field_values(cls, model: EntityModel) -> dict[str, Any]:
 
         return fields
 
+    @overload
+    def to_model(
+        self,
+        *,
+        repository_path: Optional[pathlib.Path] = None,
+        serialize_repository_content: bool = False,
+        unstored: Literal[False] = False,
+    ) -> EntityModelType: ...
+
+    @overload
+    def to_model(
+        self,
+        *,
+        repository_path: Optional[pathlib.Path] = None,
+        serialize_repository_content: bool = False,
+        unstored: Literal[True] = True,
+    ) -> BaseModel: ...
+
+    @overload
+    def to_model(
+        self,
+        *,
+        repository_path: Optional[pathlib.Path] = None,
+        serialize_repository_content: bool = False,
+        unstored: bool = False,
+    ) -> BaseModel: ...
+
     def to_model(
         self,
         *,
         repository_path: Optional[pathlib.Path] = None,
         serialize_repository_content: bool = False,
         unstored: bool = False,
-    ) -> EntityModel:
+    ) -> BaseModel:
         """Return the entity instance as an instance of its model.
 
         :param repository_path: If the orm node has files in the repository, this path is used to read the repository
@@ -318,7 +347,7 @@ def to_model(
         return Model(**fields)
 
     @classmethod
-    def from_model(cls, model: EntityModel) -> Self:
+    def from_model(cls, model: BaseModel) -> Self:
         """Return an entity instance from an instance of its model.
 
         :param model: An instance of the entity's model class.

src/aiida/orm/groups.py

@@ -20,7 +20,7 @@
 from aiida.common import exceptions
 from aiida.common.lang import classproperty, type_check
 from aiida.common.pydantic import MetadataField
-from aiida.common.typing import Self, TypeAlias
+from aiida.common.typing import Self
 from aiida.common.warnings import warn_deprecation
 from aiida.manage import get_manager
 
@@ -148,10 +148,10 @@ class GroupModel(entities.EntityModel):
     )
 
 
-class Group(entities.Entity['BackendGroup', GroupCollection]):
+class Group(entities.Entity['BackendGroup', GroupCollection, GroupModel]):
     """An AiiDA ORM implementation of group of nodes."""
 
-    Model: TypeAlias = GroupModel
+    Model = GroupModel
 
     __type_string: ClassVar[Optional[str]]
 

src/aiida/orm/logs.py

@@ -17,7 +17,6 @@
 
 from aiida.common import timezone
 from aiida.common.pydantic import MetadataField
-from aiida.common.typing import TypeAlias
 from aiida.manage import get_manager
 
 from . import entities
@@ -161,10 +160,10 @@ class LogModel(entities.EntityModel):
     )
 
 
-class Log(entities.Entity['BackendLog', LogCollection]):
+class Log(entities.Entity['BackendLog', LogCollection, LogModel]):
     """An AiiDA Log entity.  Corresponds to a logged message against a particular AiiDA node."""
 
-    Model: TypeAlias = LogModel
+    Model = LogModel
 
     _CLS_COLLECTION = LogCollection
 

src/aiida/orm/nodes/data/array/array.py

@@ -19,7 +19,6 @@
 from pydantic import field_validator
 
 from aiida.common.pydantic import MetadataField
-from aiida.common.typing import TypeAlias
 
 from .. import data
 from ..base import to_aiida_type
@@ -75,7 +74,7 @@ class ArrayData(data.Data):
 
     """
 
-    Model: TypeAlias = ArrayDataModel
+    Model = ArrayDataModel
 
     array_prefix = 'array|'
     default_array_name = 'default'

src/aiida/orm/nodes/data/array/bands.py

@@ -20,7 +20,6 @@
 
 from aiida.common.exceptions import ValidationError
 from aiida.common.pydantic import MetadataField
-from aiida.common.typing import TypeAlias
 from aiida.common.utils import join_labels, prettify_labels
 
 from . import kpoints
@@ -228,7 +227,7 @@ class BandsDataModel(kpoints.KpointsDataModel):
 class BandsData(kpoints.KpointsData):
     """Class to handle bands data"""
 
-    Model: TypeAlias = BandsDataModel
+    Model = BandsDataModel
 
     def __init__(
         self,
@@ -244,7 +243,7 @@ def set_kpointsdata(self, kpointsdata):
         """Load the kpoints from a kpoint object.
         :param kpointsdata: an instance of KpointsData class
         """
-        if not isinstance(kpointsdata, KpointsData):
+        if not isinstance(kpointsdata, kpoints.KpointsData):
             raise ValueError('kpointsdata must be of the KpointsData class')
         try:
             self.cell = kpointsdata.cell

src/aiida/orm/nodes/data/array/kpoints.py

@@ -18,7 +18,6 @@
 import numpy
 
 from aiida.common.pydantic import MetadataField
-from aiida.common.typing import TypeAlias
 
 from . import array
 
@@ -84,7 +83,7 @@ class KpointsData(array.ArrayData):
     set_cell_from_structure methods.
     """
 
-    Model: TypeAlias = KpointsDataModel
+    Model = KpointsDataModel
 
     def __init__(
         self,