Update node serialization/deserialization and other Pydantic issues

docs/source/nitpick-exceptions

@@ -26,6 +26,7 @@ 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
 
@@ -68,6 +69,7 @@ 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.BackendEntityType
 py:class aiida.orm.entities.CollectionType

environment.yml

@@ -27,7 +27,7 @@ dependencies:
 - pgsu~=0.3.0
 - psutil~=7.0
 - psycopg[binary]<4,>=3.0.2
-- pydantic~=2.6
+- pydantic~=2.12
 - pytz~=2021.1
 - pyyaml~=6.0
 - requests~=2.0

pyproject.toml

@@ -50,7 +50,7 @@ dependencies = [
   'pgsu~=0.3.0',
   'psutil~=7.0',
   'psycopg[binary]>=3.0.2,<4',
-  'pydantic~=2.6',
+  'pydantic~=2.12',
   'pytz~=2021.1',
   'pyyaml~=6.0',
   'requests~=2.0',

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
 
@@ -25,17 +27,22 @@
 from aiida.cmdline.utils.common import validate_output_filename
 from aiida.cmdline.utils.decorators import with_dbenv
 from aiida.common import exceptions
+from aiida.common.pydantic import get_metadata
+
+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)
     except (TypeError, ValueError) as exception:
         echo.echo_critical(f'Failed to create instance `{cls}`: {exception}')
 
@@ -223,50 +230,27 @@ def code_duplicate(ctx, code, non_interactive, **kwargs):
 @verdi_code.command()
 @arguments.CODE()
 @with_dbenv()
-def show(code):
+def show(code: Code):
     """Display detailed information for a code."""
     from aiida.cmdline import is_verbose
-    from aiida.common.pydantic import get_metadata
 
     table = []
 
     # These are excluded from the CLI, so we add them manually
     table.append(['PK', code.pk])
     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,
-        ):
-            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,
-        ):
-            continue
-
-        value = getattr(code, field_name)
+    table.append(['Type', code.entry_point.name if code.entry_point else None])
+    table.append(['Label', code.label])
+    table.append(['Description', code.description or ''])
 
-        # Special handling for computer field to show additional info
-        if field_name == 'computer':
-            value = f'{value.label} ({value.hostname}), pk: {value.pk}'
+    if code.computer is not None:
+        table.append(['Computer', f'{code.computer.label} ({code.computer.hostname}), pk: {code.computer.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])
+    for key, field in code.AttributesModel.model_fields.items():
+        if key == 'source' or get_metadata(field, 'write_only'):
+            continue
+        value = getattr(code, key)
+        table.append([field.title, value])
 
     if is_verbose():
         table.append(['Calculations', len(code.base.links.get_outgoing().all())])

src/aiida/cmdline/groups/dynamic.py

@@ -9,6 +9,7 @@
 import click
 
 from aiida.common import exceptions
+from aiida.common.pydantic import OrmModel
 from aiida.plugins.entry_point import ENTRY_POINT_GROUP_FACTORY_MAPPING, get_entry_point_names
 from aiida.plugins.factories import BaseFactory
 
@@ -96,9 +97,15 @@ def call_command(self, ctx: click.Context, cls: t.Any, non_interactive: bool, **
         from pydantic import ValidationError
 
         if hasattr(cls, 'Model'):
-            # The plugin defines a pydantic model: use it to validate the provided arguments
             try:
-                cls.Model(**kwargs)
+                Model = getattr(cls, 'CreateModel', cls.Model)  # noqa: N806
+                if 'attributes' in Model.model_fields:
+                    attr_field = Model.model_fields['attributes']
+                    AttributesModel = t.cast(OrmModel, attr_field.annotation)  # noqa: N806
+                    kwargs['attributes'] = {
+                        key: kwargs.pop(key) for key in AttributesModel.model_fields.keys() if key in kwargs
+                    }
+                Model(**kwargs)
             except ValidationError as exception:
                 param_hint = [
                     f'--{loc.replace("_", "-")}'  # type: ignore[union-attr]
@@ -153,8 +160,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 +173,18 @@ 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'):
+        fields = dict(Model.model_fields)
+        attr_field = fields.get('attributes')
+        if attr_field is not None:
+            AttributesModel = t.cast(OrmModel, attr_field.annotation)  # noqa: N806
+            fields |= AttributesModel.model_fields
+
+        for key, field_info in fields.items():
+            if key in ('extras', 'attributes'):
                 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,18 +3,16 @@
 from __future__ import annotations
 
 import typing as t
-from pathlib import Path
 
-from pydantic import Field
+from pydantic import BaseModel, ConfigDict, Field, create_model
+from pydantic.fields import FieldInfo
 from pydantic_core import PydanticUndefined
 
 if t.TYPE_CHECKING:
-    from pydantic import BaseModel
-
     from aiida.orm import Entity
 
 
-def get_metadata(field_info: t.Any, key: str, default: t.Any | None = None) -> t.Any:
+def get_metadata(field_info: FieldInfo, key: str, default: t.Any | None = None) -> t.Any:
     """Return a the metadata of the given field for a particular key.
 
     :param field_info: The field from which to retrieve the metadata.
@@ -23,31 +21,109 @@ def get_metadata(field_info: t.Any, key: str, default: t.Any | None = None) -> t
     :returns: The metadata if defined, otherwise the default.
     """
     for element in field_info.metadata:
-        if key in element:
+        if isinstance(element, dict) and key in element:
             return element[key]
     return default
 
 
+class OrmModel(BaseModel, defer_build=True):
+    """Base class for all ORM entity models."""
+
+    model_config = ConfigDict(extra='forbid')
+
+    @classmethod
+    def __pydantic_init_subclass__(cls, **kwargs: t.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_create_model(cls: t.Type[OrmModel]) -> t.Type[OrmModel]:
+        """Return a derived creation model class with read-only fields removed.
+
+        This also removes any serializers/validators defined on those fields.
+
+        :return: The derived creation model class.
+        """
+
+        cached = cls.__dict__.get('_AIIDA_CREATE_MODEL')
+        if isinstance(cached, type) and issubclass(cached, OrmModel):
+            return cached
+
+        CreateModel = create_model(  # noqa: N806
+            'CreateModel',
+            __base__=cls,
+            __module__=cls.__module__,
+            __qualname__=cls.__qualname__.replace('Model', 'CreateModel'),
+        )
+        CreateModel.model_config['extra'] = 'ignore'
+        CreateModel.model_config['json_schema_extra'] = {
+            **CreateModel.model_config.get('json_schema_extra', {}),  # type: ignore[dict-item]
+            'additionalProperties': False,
+        }
+
+        CreateModel.model_rebuild(force=True)
+
+        readonly_fields: t.List[str] = []
+        for key, field in CreateModel.model_fields.items():
+            annotation = field.annotation
+            if get_metadata(field, 'read_only'):
+                readonly_fields.append(key)
+            elif isinstance(annotation, type) and issubclass(annotation, OrmModel):
+                sub_create_model = annotation._as_create_model()
+                field.annotation = sub_create_model
+                if any(f.is_required() for f in sub_create_model.model_fields.values()):
+                    field.default_factory = None
+
+        # Remove read-only fields
+        for key in readonly_fields:
+            CreateModel.model_fields.pop(key, None)
+            if hasattr(CreateModel, key):
+                delattr(CreateModel, key)
+
+        # Prune field validators/serializers referring to read-only fields
+        decorators = CreateModel.__pydantic_decorators__
+
+        def prune_field_decorators(field_decorators: dict[str, t.Any]) -> dict[str, t.Any]:
+            return {
+                key: decorator
+                for key, decorator in field_decorators.items()
+                if all(field not in readonly_fields for field in decorator.info.fields)
+            }
+
+        decorators.field_validators = prune_field_decorators(decorators.field_validators)
+        decorators.field_serializers = prune_field_decorators(decorators.field_serializers)
+
+        # Make subsequent calls idempotent for this specific class
+        cls._AIIDA_CREATE_MODEL = CreateModel  # type: ignore[attr-defined]
+
+        CreateModel.model_rebuild(force=True)
+        return CreateModel
+
+
 def MetadataField(  # noqa: N802
     default: t.Any = PydanticUndefined,
     *,
     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,
-    exclude_to_orm: bool = False,
-    exclude_from_cli: bool = False,
-    is_attribute: bool = True,
-    is_subscriptable: bool = False,
+    orm_to_model: t.Callable[[Entity[t.Any, t.Any], dict[str, t.Any]], t.Any] | None = None,
+    model_to_orm: t.Callable[[OrmModel], t.Any] | None = None,
+    read_only: bool = False,
+    write_only: bool = False,
+    may_be_large: bool = False,
     **kwargs: t.Any,
 ) -> t.Any:
     """Return a :class:`pydantic.fields.Field` instance with additional metadata.
 
     .. code-block:: python
 
-        class Model(BaseModel):
+        class Model(OrmModel):
 
             attribute: MetadataField('default', priority=1000, short_name='-A')
 
@@ -66,20 +142,34 @@ 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
         primary key.
     :param orm_to_model: Optional callable to convert the value of a field from an ORM instance to a model instance.
     :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
+    :param read_only: 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 write_only: When set to ``True``, this field value will not be populated when constructing the model from an
+        ORM entity through ``Entity.to_model``.
+    :param may_be_large: Whether the field value may be large. This is used to determine whether to include the field
+        when serializing the entity for various purposes, such as exporting or logging.
     """
+
+    extra = kwargs.pop('json_schema_extra', {})
+
+    if read_only and write_only:
+        raise ValueError('A field cannot be both read-only and write-only.')
+
+    if read_only:
+        extra.update({'readOnly': True})
+
+    if write_only:
+        extra.update({'writeOnly': True})
+
+    kwargs['json_schema_extra'] = extra
+
     field_info = Field(default, **kwargs)
 
     for key, value in (
@@ -89,10 +179,9 @@ class Model(BaseModel):
         ('orm_class', orm_class),
         ('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),
+        ('read_only', read_only),
+        ('write_only', write_only),
+        ('may_be_large', may_be_large),
     ):
         if value is not None:
             field_info.metadata.append({key: value})