Add Model class property on Entity

Author: edan-bainglass

src/aiida/cmdline/groups/dynamic.py

@@ -95,7 +95,7 @@ def call_command(self, ctx: click.Context, cls: t.Any, non_interactive: bool, **
         """Call the ``command`` after validating the provided inputs."""
         from pydantic import ValidationError
 
-        if hasattr(cls, 'Model') or hasattr(cls, '_CLS_MODEL'):
+        if hasattr(cls, 'Model'):
             # The plugin defines a pydantic model: use it to validate the provided arguments
             Model = cls.InputModel if hasattr(cls, 'InputModel') else cls.Model  # noqa: N806
             try:
@@ -158,7 +158,7 @@ def list_options(self, entry_point: str) -> list[t.Callable[[FC], FC]]:
 
         cls = self.factory(entry_point)
 
-        if not (hasattr(cls, 'Model') or hasattr(cls, '_CLS_MODEL')):
+        if not hasattr(cls, 'Model'):
             from aiida.common.warnings import warn_deprecation
 
             warn_deprecation(

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

@@ -29,7 +29,7 @@
 __all__ = ('AuthInfo',)
 
 
-class AuthInfoCollection(entities.Collection['AuthInfo']):
+class AuthInfoCollection(entities.EntityCollection['AuthInfo']):
     """The collection of `AuthInfo` entries."""
 
     @staticmethod
@@ -44,7 +44,7 @@ def delete(self, pk: int) -> None:
         self._backend.authinfos.delete(pk)
 
 
-class Model(entities.Model):
+class AuthInfoModel(entities.EntityModel):
     computer: int = MetadataField(
         description='The PK of the computer',
         is_attribute=False,
@@ -74,11 +74,11 @@ class Model(entities.Model):
     )
 
 
-class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection, Model]):
+class AuthInfo(entities.Entity['BackendAuthInfo', AuthInfoCollection, AuthInfoModel]):
     """ORM class that models the authorization information that allows a `User` to connect to a `Computer`."""
 
     _CLS_COLLECTION = AuthInfoCollection
-    _CLS_MODEL = Model
+    _CLS_MODEL = AuthInfoModel
 
     PROPERTY_WORKDIR = 'workdir'
 

src/aiida/orm/comments.py

@@ -28,7 +28,7 @@
 __all__ = ('Comment',)
 
 
-class CommentCollection(entities.Collection['Comment']):
+class CommentCollection(entities.EntityCollection['Comment']):
     """The collection of Comment entries."""
 
     @staticmethod
@@ -65,7 +65,7 @@ def delete_many(self, filters: dict) -> List[int]:
         return self._backend.comments.delete_many(filters)
 
 
-class Model(entities.Model):
+class CommentModel(entities.EntityModel):
     uuid: UUID = MetadataField(
         description='The UUID of the comment',
         is_attribute=False,
@@ -99,11 +99,11 @@ class Model(entities.Model):
     )
 
 
-class Comment(entities.Entity['BackendComment', CommentCollection, Model]):
+class Comment(entities.Entity['BackendComment', CommentCollection, CommentModel]):
     """Base class to map a DbComment that represents a comment attached to a certain Node."""
 
     _CLS_COLLECTION = CommentCollection
-    _CLS_MODEL = Model
+    _CLS_MODEL = CommentModel
 
     def __init__(
         self, node: 'Node', user: 'User', content: Optional[str] = None, backend: Optional['StorageBackend'] = None

src/aiida/orm/computers.py

@@ -31,7 +31,7 @@
 __all__ = ('Computer',)
 
 
-class ComputerCollection(entities.Collection['Computer']):
+class ComputerCollection(entities.EntityCollection['Computer']):
     """The collection of Computer entries."""
 
     @staticmethod
@@ -64,7 +64,7 @@ def delete(self, pk: int) -> None:
         return self._backend.computers.delete(pk)
 
 
-class Model(entities.Model):
+class ComputerModel(entities.EntityModel):
     uuid: UUID = MetadataField(
         description='The UUID of the computer',
         is_attribute=False,
@@ -98,7 +98,7 @@ class Model(entities.Model):
     )
 
 
-class Computer(entities.Entity['BackendComputer', ComputerCollection, Model]):
+class Computer(entities.Entity['BackendComputer', ComputerCollection, ComputerModel]):
     """Computer entity."""
 
     _logger = AIIDA_LOGGER.getChild('orm.computers')
@@ -109,7 +109,7 @@ class Computer(entities.Entity['BackendComputer', ComputerCollection, Model]):
     PROPERTY_SHEBANG = 'shebang'
 
     _CLS_COLLECTION = ComputerCollection
-    _CLS_MODEL = Model
+    _CLS_MODEL = ComputerModel
 
     def __init__(
         self,

src/aiida/orm/entities.py

@@ -35,14 +35,14 @@
     from aiida.orm.querybuilder import FilterType, OrderByType, QueryBuilder
 
 __all__ = (
-    'Collection',
     'Entity',
+    'EntityCollection',
+    'EntityModel',
     'EntityTypes',
-    'Model',
 )
 
-CollectionType = TypeVar('CollectionType', bound='Collection[Any]')
-EntityModelType = TypeVar('EntityModelType', bound='Model')
+CollectionType = TypeVar('CollectionType', bound='EntityCollection[Any]')
+EntityModelType = TypeVar('EntityModelType', bound='EntityModel')
 EntityType = TypeVar('EntityType', bound='Entity[Any,Any,Any]')
 BackendEntityType = TypeVar('BackendEntityType', bound='BackendEntity')
 
@@ -61,7 +61,7 @@ class EntityTypes(Enum):
     GROUP_NODE = 'group_node'
 
 
-class Collection(abc.ABC, Generic[EntityType]):
+class EntityCollection(abc.ABC, Generic[EntityType]):
     """Container class that represents the collection of objects of a particular entity type."""
 
     @staticmethod
@@ -184,7 +184,7 @@ def count(self, filters: Optional['FilterType'] = None) -> int:
         return self.query(filters=filters).count()
 
 
-class Model(BaseModel, defer_build=True):
+class EntityModel(BaseModel, defer_build=True):
     model_config = ConfigDict(extra='forbid')
 
     pk: int = MetadataField(
@@ -194,16 +194,6 @@ class Model(BaseModel, defer_build=True):
         exclude_from_cli=True,
     )
 
-    @classmethod
-    def __pydantic_init_subclass__(cls, **kwargs: Any) -> None:
-        """Sets the JSON schema title of the model.
-
-        The qualified name of the class is used, with dots removed. For example, `Node.Model` becomes `NodeModel`
-        in the JSON schema.
-        """
-        super().__pydantic_init_subclass__(**kwargs)
-        cls.model_config['title'] = cls.__qualname__.replace('.', '')
-
     @classmethod
     def as_input_model(cls: Type[EntityModelType]) -> Type[EntityModelType]:
         """Return a derived model class with read-only fields removed.
@@ -214,7 +204,7 @@ def as_input_model(cls: Type[EntityModelType]) -> Type[EntityModelType]:
         """
 
         # Derive the input model from the original model
-        new_name = cls.__qualname__.replace('.Model', 'InputModel')
+        new_name = cls.__qualname__.replace('Model', 'InputModel')
         InputModel = create_model(  # noqa: N806
             new_name,
             __base__=cls,
@@ -257,10 +247,18 @@ def _prune_field_decorators(field_decorators: dict[str, Any]) -> dict[str, Any]:
 class Entity(abc.ABC, Generic[BackendEntityType, CollectionType, EntityModelType], metaclass=EntityFieldMeta):
     """An AiiDA entity"""
 
-    _CLS_COLLECTION: type[CollectionType] = Collection  # type: ignore[assignment]
-    _CLS_MODEL: type[EntityModelType] = Model  # type: ignore[assignment]
+    _CLS_COLLECTION: type[CollectionType] = EntityCollection  # type: ignore[assignment]
+    _CLS_MODEL: type[EntityModelType] = EntityModel  # type: ignore[assignment]
     _logger = log.AIIDA_LOGGER.getChild('orm.entities')
 
+    @classproperty
+    def Model(cls) -> Type[EntityModelType]:  # noqa: N802, N805
+        """Return the model class for this entity.
+
+        :return: The model class.
+        """
+        return cls._CLS_MODEL
+
     @classproperty
     def InputModel(cls) -> Type[EntityModelType]:  # noqa: N802, N805
         """Return the input version of the model class for this entity.
@@ -272,9 +270,7 @@ def InputModel(cls) -> Type[EntityModelType]:  # noqa: N802, N805
     @classmethod
     def model_to_orm_fields(cls) -> dict[str, FieldInfo]:
         return {
-            key: field
-            for key, field in cls._CLS_MODEL.model_fields.items()
-            if not get_metadata(field, 'exclude_to_orm')
+            key: field for key, field in cls.Model.model_fields.items() if not get_metadata(field, 'exclude_to_orm')
         }
 
     @classmethod
@@ -325,7 +321,7 @@ def to_model(
             with `exclude_to_orm=True`.
         :return: An instance of the entity's model class.
         """
-        Model = self.InputModel if unstored else self._CLS_MODEL  # noqa: N806
+        Model = self.InputModel if unstored else self.Model  # noqa: N806
         fields = self._collect_model_field_values(
             repository_path=repository_path,
             serialize_repository_content=serialize_repository_content,
@@ -396,7 +392,7 @@ def from_serialized(cls, serialized: dict[str, Any], unstored: bool = False) ->
         cls._logger.warning(
             'Serialization through pydantic is still an experimental feature and might break in future releases.'
         )
-        Model = cls.InputModel if unstored else cls._CLS_MODEL  # noqa: N806
+        Model = cls.InputModel if unstored else cls.Model  # noqa: N806
         return cls.from_model(Model(**serialized))
 
     @classproperty
@@ -542,7 +538,7 @@ def _collect_model_field_values(
         """
         fields: dict[str, Any] = {}
 
-        Model = self.InputModel if unstored else self._CLS_MODEL  # noqa: N806
+        Model = self.InputModel if unstored else self.Model  # noqa: N806
 
         for key, field in Model.model_fields.items():
             if skip_cli_excluded and get_metadata(field, 'exclude_from_cli'):

src/aiida/orm/fields.py

@@ -410,65 +410,65 @@ def __init__(cls, name, bases, classdict):
 
         fields = {}
 
-        # If the class has an attribute ``_CLS_MODEL`` that is a subclass of :class:`pydantic.BaseModel`, parse the
-        # model fields to build up the ``fields`` class attribute, which is used to allow specifying ``QueryBuilder``
-        # filters programmatically.
-        if hasattr(cls, '_CLS_MODEL') and issubclass(cls._CLS_MODEL, BaseModel):
-            # If the class itself directly specifies the ``_CLS_MODEL`` attribute, check that it is valid. Here, the
-            # check ``cls.__dict__`` is used instead of ``hasattr`` as the former only returns true if the class itself
+        # If the class has an attribute ``Model`` that is a subclass of :class:`pydantic.BaseModel`, parse the model
+        # fields to build up the ``fields`` class attribute, which is used to allow specifying ``QueryBuilder`` filters
+        # programmatically.
+        if hasattr(cls, 'Model') and issubclass(cls.Model, BaseModel):
+            # If the class itself directly specifies the ``Model`` attribute, check that it is valid. Here, the check
+            # ``cls.__dict__`` is used instead of ``hasattr`` as the former only returns true if the class itself
             # defines the attribute and does not just inherit it from a base class. In that case, this check will
             # already have been performed for that subclass.
 
-            # When a class defines a ``_CLS_MODEL``, the following check ensures that the model inherits from the same
-            # bases as the class containing the attribute itself. For example, if ``cls`` inherits from ``ClassA`` and
-            # ``ClassB`` that each define a ``_CLS_MODEL``, the ``cls._CLS_MODEL`` class should inherit from both
-            # ``ClassA._CLS_MODEL`` and ``ClassB._CLS_MODEL`` or it will be losing the attributes of some of the models.
-            if '_CLS_MODEL' in cls.__dict__:
-                # Get all the base classes in the MRO of this class that define a class attribute ``_CLS_MODEL`` that
-                # is a subclass of pydantic's ``BaseModel`` and not the class itself.
+            # When a class defines a ``Model``, the following check ensures that the model inherits from the same bases
+            # as the class containing the attribute itself. For example, if ``cls`` inherits from ``ClassA`` and
+            # ``ClassB`` that each define a ``Model``, the ``cls.Model`` class should inherit from both ``ClassA.Model``
+            # and ``ClassBModel`` or it will be losing the attributes of some of the models.
+            if 'Model' in cls.__dict__:
+                # Get all the base classes in the MRO of this class that define a class attribute ``Model`` that is a
+                # subclass of pydantic's ``BaseModel`` and not the class itself
                 cls_bases_with_model = [
                     base
                     for base in cls.__mro__
-                    if base is not cls and '_CLS_MODEL' in base.__dict__ and issubclass(base._CLS_MODEL, BaseModel)  # type: ignore[attr-defined]
+                    if base is not cls and 'Model' in base.__dict__ and issubclass(base.Model, BaseModel)  # type: ignore[attr-defined]
                 ]
 
                 # Now get the "leaf" bases, i.e., those base classes in the subclass list that themselves do not have a
-                # subclass in the tree. This set should be the base classes for the class' ``_CLS_MODEL`` attribute.
+                # subclass in the tree. This set should be the base classes for the class' ``Model`` attribute.
                 cls_bases_with_model_leaves = {
                     base
                     for base in cls_bases_with_model
                     if all(
-                        not issubclass(b._CLS_MODEL, base._CLS_MODEL)  # type: ignore[attr-defined]
+                        not issubclass(b.Model, base.Model)  # type: ignore[attr-defined]
                         for b in cls_bases_with_model
                         if b is not base
                     )
                 }
 
-                cls_model_bases = {base._CLS_MODEL for base in cls_bases_with_model_leaves}  # type: ignore[attr-defined]
+                cls_model_bases = {base.Model for base in cls_bases_with_model_leaves}  # type: ignore[attr-defined]
 
-                # If the base class does not have a base that defines a model, it means the ``_CLS_MODEL`` should simply
-                # have ``pydantic.BaseModel`` as its sole base.
+                # If the base class does not have a base that defines a model, it means the ``Model`` should simply have
+                # ``pydantic.BaseModel`` as its sole base.
                 if not cls_model_bases:
                     cls_model_bases = {
                         BaseModel,
                     }
 
-                # Get the set of bases of ``cls._CLS_MODEL`` that are a subclass of :class:`pydantic.BaseModel`.
-                model_bases = {base for base in cls._CLS_MODEL.__bases__ if issubclass(base, BaseModel)}
+                # Get the set of bases of ``cls.Model`` that are a subclass of :class:`pydantic.BaseModel`.
+                model_bases = {base for base in cls.Model.__bases__ if issubclass(base, BaseModel)}
 
-                # For ``cls._CLS_MODEL`` to be valid, the bases that contain a model, should equal to the leaf bases of
-                # the ``cls`` itself that also define a model.
+                # For ``cls.Model`` to be valid, the bases that contain a model, should equal to the leaf bases of the
+                # ``cls`` itself that also define a model.
                 if model_bases != cls_model_bases and not getattr(cls, '_SKIP_MODEL_INHERITANCE_CHECK', False):
-                    bases = [f'{e.__module__}.Model' for e in cls_bases_with_model_leaves]
+                    bases = [f'{e.__module__}.{e.__name__}.Model' for e in cls_bases_with_model_leaves]
                     raise RuntimeError(
-                        f'`{cls.__module__}.Model` does not subclass all necessary base classes. It should be: '
+                        f'`{cls.__name__}.Model` does not subclass all necessary base classes. It should be: '
                         f'`class Model({", ".join(sorted(bases))}):`'
                     )
 
-            for key, field in cls._CLS_MODEL.model_fields.items():
+            for key, field in cls.Model.model_fields.items():
                 fields[key] = add_field(
                     key,
-                    alias=field.alias,
+                    alias=get_metadata(field, 'alias', None),
                     dtype=field.annotation,
                     doc=field.description,
                     is_attribute=get_metadata(field, 'is_attribute', False),

src/aiida/orm/groups.py

@@ -57,7 +57,7 @@ def load_group_class(type_string: str) -> Type[Group]:
     return group_class
 
 
-class GroupCollection(entities.Collection['Group']):
+class GroupCollection(entities.EntityCollection['Group']):
     """Collection of Groups"""
 
     @staticmethod
@@ -107,7 +107,7 @@ def extras(self) -> extras.EntityExtras:
         return extras.EntityExtras(self._group)
 
 
-class Model(entities.Model):
+class GroupModel(entities.EntityModel):
     uuid: UUID = MetadataField(
         description='The UUID of the group',
         is_attribute=False,
@@ -148,13 +148,13 @@ class Model(entities.Model):
     )
 
 
-class Group(entities.Entity['BackendGroup', GroupCollection, Model]):
+class Group(entities.Entity['BackendGroup', GroupCollection, GroupModel]):
     """An AiiDA ORM implementation of group of nodes."""
 
     __type_string: ClassVar[Optional[str]]
 
     _CLS_COLLECTION = GroupCollection
-    _CLS_MODEL = Model
+    _CLS_MODEL = GroupModel
 
     def __init__(
         self,