Extract Model classes out of ORM entity classes

docs/source/nitpick-exceptions

@@ -26,14 +26,16 @@ py:class json.encoder.JSONEncoder
 py:class EXPOSED_TYPE
 py:class EVENT_CALLBACK_TYPE
 py:class datetime
+py:class UUID
 py:class types.LambdaType
 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
@@ -50,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
@@ -61,21 +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.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
@@ -88,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

@@ -331,7 +331,7 @@ enable_error_code = [
   'truthy-bool'
 ]
 plugins = [
-  'pydantic.mypy',
+  # 'pydantic.mypy',
   'sqlalchemy.ext.mypy.plugin'
 ]
 scripts_are_modules = true
@@ -480,7 +480,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/cmdline/commands/cmd_code.py

@@ -8,11 +8,13 @@
 ###########################################################################
 """`verdi code` command."""
 
+from __future__ import annotations
+
 import pathlib
 import warnings
 from collections import defaultdict
 from functools import partial
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import click
 
@@ -26,16 +28,20 @@
 from aiida.cmdline.utils.decorators import with_dbenv
 from aiida.common import exceptions
 
+if TYPE_CHECKING:
+    from aiida.orm import Code
+
 
 @verdi.group('code')
 def verdi_code():
     """Setup and manage codes."""
 
 
-def create_code(ctx: click.Context, cls, **kwargs) -> None:
+def create_code(ctx: click.Context, cls: Code, **kwargs) -> None:
     """Create a new `Code` instance."""
     try:
-        instance = cls._from_model(cls.Model(**kwargs))
+        model = cls.CreateModel(**kwargs)
+        instance = cls.from_model(model)  # type: ignore[arg-type]
     except (TypeError, ValueError) as exception:
         echo.echo_critical(f'Failed to create instance `{cls}`: {exception}')
 
@@ -226,7 +232,6 @@ def code_duplicate(ctx, code, non_interactive, **kwargs):
 def show(code):
     """Display detailed information for a code."""
     from aiida.cmdline import is_verbose
-    from aiida.common.pydantic import get_metadata
 
     table = []
 
@@ -235,38 +240,21 @@ def show(code):
     table.append(['UUID', code.uuid])
     table.append(['Type', code.entry_point.name])
 
-    for field_name, field_info in code.Model.model_fields.items():
-        # Skip fields excluded from CLI
-        if get_metadata(
-            field_info,
-            key='exclude_from_cli',
-            default=False,
-        ):
+    for field_name, field_info in code.CreateModel.model_fields.items():
+        # We don't show extras for codes
+        if field_name == 'extras':
             continue
 
-        # Skip fields that are not stored in the attributes column
-        # NOTE: this also catches e.g., filepath_files for PortableCode, which is actually a "misuse"
-        # of the is_attribute metadata flag, as there it is flagging that the field is not stored at all!
-        # TODO (edan-bainglass) consider improving this by introducing a new metadata flag or reworking PortableCode
-        # TODO see also Dict and InstalledCode for other potential misuses of is_attribute
-        if not get_metadata(
-            field_info,
-            key='is_attribute',
-            default=True,
-        ):
+        # FIXME resolve this hardcoded special case properly
+        if field_name == 'filepath_files':
             continue
 
         value = getattr(code, field_name)
 
-        # Special handling for computer field to show additional info
         if field_name == 'computer':
             value = f'{value.label} ({value.hostname}), pk: {value.pk}'
 
-        # Use the field's title as display name.
-        # This allows for custom titles (class-cased by default from Pydantic).
-        display_name = field_info.title
-
-        table.append([display_name, value])
+        table.append([field_info.title, value])
 
     if is_verbose():
         table.append(['Calculations', len(code.base.links.get_outgoing().all())])

src/aiida/cmdline/groups/dynamic.py

@@ -97,8 +97,9 @@ def call_command(self, ctx: click.Context, cls: t.Any, non_interactive: bool, **
 
         if hasattr(cls, 'Model'):
             # The plugin defines a pydantic model: use it to validate the provided arguments
+            Model = getattr(cls, 'CreateModel', cls.Model)  # noqa: N806
             try:
-                cls.Model(**kwargs)
+                Model(**kwargs)
             except ValidationError as exception:
                 param_hint = [
                     f'--{loc.replace("_", "-")}'  # type: ignore[union-attr]
@@ -153,8 +154,6 @@ def list_options(self, entry_point: str) -> list[t.Callable[[FC], FC]]:
         """
         from pydantic_core import PydanticUndefined
 
-        from aiida.common.pydantic import get_metadata
-
         cls = self.factory(entry_point)
 
         if not hasattr(cls, 'Model'):
@@ -168,10 +167,13 @@ def list_options(self, entry_point: str) -> list[t.Callable[[FC], FC]]:
             options_spec = self.factory(entry_point).get_cli_options()  # type: ignore[union-attr]
             return [self.create_option(*item) for item in options_spec]
 
+        Model = getattr(cls, 'CreateModel', cls.Model)  # noqa: N806
+
         options_spec = {}
 
-        for key, field_info in cls.Model.model_fields.items():
-            if get_metadata(field_info, 'exclude_from_cli'):
+        for key, field_info in Model.model_fields.items():
+            # We do not prompt for extras
+            if key == 'extras':
                 continue
 
             default = field_info.default_factory if field_info.default is PydanticUndefined else field_info.default

src/aiida/common/datastructures.py

@@ -167,8 +167,8 @@ class CalcInfo(DefaultFieldsAttributeDict):
         max_wallclock_seconds: None | int
         max_memory_kb: None | int
         rerunnable: bool
-        retrieve_list: None | list[str | tuple[str, str, str]]
-        retrieve_temporary_list: None | list[str | tuple[str, str, str]]
+        retrieve_list: None | list[str | tuple[str, str, int]]
+        retrieve_temporary_list: None | list[str | tuple[str, str, int]]
         local_copy_list: None | list[tuple[str, str, str]]
         remote_copy_list: None | list[tuple[str, str, str]]
         remote_symlink_list: None | list[tuple[str, str, str]]

src/aiida/common/pydantic.py

@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import typing as t
-from pathlib import Path
 
 from pydantic import Field
 from pydantic_core import PydanticUndefined
@@ -34,11 +33,10 @@ 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], Path], t.Any] | None = None,
-    model_to_orm: t.Callable[['BaseModel'], t.Any] | None = None,
+    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], dict[str, t.Any]], t.Any] | None = None,
+    model_to_orm: t.Callable[[BaseModel], t.Any] | None = None,
     exclude_to_orm: bool = False,
-    exclude_from_cli: bool = False,
     is_attribute: bool = True,
     is_subscriptable: bool = False,
     **kwargs: t.Any,
@@ -66,7 +64,7 @@ class Model(BaseModel):
     :param short_name: Optional short name to use for an option on a command line interface.
     :param option_cls: The :class:`click.Option` class to use to construct the option.
     :param orm_class: The class, or entry point name thereof, to which the field should be converted. If this field is
-        defined, the value of this field should acccept an integer which will automatically be converted to an instance
+        defined, the value of this field should accept an integer which will automatically be converted to an instance
         of said ORM class using ``orm_class.collection.get(id={field_value})``. This is useful, for example, where a
         field represents an instance of a different entity, such as an instance of ``User``. The serialized data would
         store the ``pk`` of the user, but the ORM entity instance would receive the actual ``User`` instance with that
@@ -75,11 +73,14 @@ class Model(BaseModel):
     :param model_to_orm: Optional callable to convert the value of a field from a model instance to an ORM instance.
     :param exclude_to_orm: When set to ``True``, this field value will not be passed to the ORM entity constructor
         through ``Entity.from_model``.
-    :param exclude_from_cli: When set to ``True``, this field value will not be exposed on the CLI command that is
-        dynamically generated to create a new instance.
-    :param is_attribute: Whether the field is stored as an attribute.
-    :param is_subscriptable: Whether the field can be indexed like a list or dictionary.
+    :param is_attribute: Whether the field is stored as an attribute. Used by `QbFields`.
+    :param is_subscriptable: Whether the field can be indexed like a list or dictionary. Used by `QbFields`.
     """
+    if exclude_to_orm:
+        extra = kwargs.pop('json_schema_extra', {})
+        extra.update({'readOnly': True})
+        kwargs['json_schema_extra'] = extra
+
     field_info = Field(default, **kwargs)
 
     for key, value in (
@@ -90,7 +91,6 @@ class Model(BaseModel):
         ('orm_to_model', orm_to_model),
         ('model_to_orm', model_to_orm),
         ('exclude_to_orm', exclude_to_orm),
-        ('exclude_from_cli', exclude_from_cli),
         ('is_attribute', is_attribute),
         ('is_subscriptable', is_subscriptable),
     ):

src/aiida/orm/__init__.py

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