✨ Store a modelled socket through its own model

src/aiida_workgraph/decorator.py

@@ -1,16 +1,21 @@
 from __future__ import annotations
-from typing import Callable, Dict, Optional, Union
+from dataclasses import replace
+from typing import Callable, Dict, Optional, Type, Union
+from pydantic import BaseModel
 from aiida.engine import calcfunction, workfunction, CalcJob, WorkChain
 from aiida_workgraph.task import Task
 from .workgraph import WorkGraph
 import inspect
 from .task import TaskHandle
 from node_graph.task_spec import TaskSpec
 from node_graph.socket_spec import SocketSpec
+from aiida_workgraph.socket_spec import SocketSpecAPI, node_typed_paths
 from aiida_workgraph.tasks.aiida import _build_aiida_function_taskspec
 from node_graph.error_handler import ErrorHandlerSpec, normalize_error_handlers
 from aiida_workgraph.tasks.pythonjob_tasks import build_pyfunction_taskspec
 from aiida_workgraph.tasks.aiida import AiiDAProcessTask
+from node_graph.executor import RuntimeExecutor
+from node_graph.input_model import ModelContractError, apply_models, rebind_executor_callable
 
 
 def _spec_for(
@@ -120,6 +125,52 @@ def decorator_task_wrapper(*args, **kwargs):
     return decorator_task_wrapper
 
 
+def _refuse_node_typed_inputs(model: Optional[Type[BaseModel]], spec: Optional[SocketSpec]) -> None:
+    """Raise when a model asks a PyFunction body for a value it cannot be handed.
+
+    A PyFunction's inputs are read out of their nodes before its body runs, so
+    a field declaring an AiiDA type would be handed what the node carries and
+    the model would refuse the very thing it asked for. A calcfunction's body
+    is handed the nodes themselves, and is where such a field belongs.
+    """
+    if model is None or spec is None:
+        return
+    paths = node_typed_paths(spec)
+    if not paths:
+        return
+    listed = ', '.join(repr(path) for path in paths)
+    raise ModelContractError(
+        f'{model.__name__} declares {listed} as an AiiDA type, and a task declared with '
+        '@task runs its body as a PyFunction, which is handed the value a node carries, '
+        'never the node.\n'
+        'How to fix: declare the task with @task.calcfunction, whose body is handed the '
+        'node; or declare the field as the Python type the body reads.'
+    )
+
+
+def _refuse_namespace_inputs(model: Optional[Type[BaseModel]], spec: Optional[SocketSpec]) -> None:
+    """Raise when a model gives a calcfunction a parameter AiiDA cannot express.
+
+    A process function's parameter is one port carrying one node. A field
+    declaring a nested model, or a ``dict[str, T]``, asks for a namespace, and
+    AiiDA turns the mapping it is handed into a single ``orm.Dict``: the
+    members lose the nodes they were, and the model refuses the ``Dict`` it
+    never declared.
+    """
+    if model is None or spec is None:
+        return
+    namespaces = [name for name, field in (spec.fields or {}).items() if field.is_namespace()]
+    if not namespaces:
+        return
+    listed = ', '.join(repr(name) for name in namespaces)
+    raise ModelContractError(
+        f'{model.__name__} declares {listed} as a namespace -- a nested model or a '
+        'dict[str, T] -- and a calcfunction parameter is one port carrying one node.\n'
+        'How to fix: declare the task with @task, whose body is handed a namespace as a '
+        'mapping; or give the model one field per value the body reads.'
+    )
+
+
 class TaskDecoratorCollection:
     """Collection of task decorators."""
 
@@ -131,6 +182,8 @@ def decorator_task(
         outputs: Optional[SocketSpec | list] = None,
         error_handlers: Optional[Dict[str, ErrorHandlerSpec]] = None,
         catalog: str = 'Others',
+        input_model: Optional[Type[BaseModel]] = None,
+        output_model: Optional[Type[BaseModel]] = None,
     ) -> Callable:
         """Generate a decorator that register a function as a task.
 
@@ -139,21 +192,34 @@ def decorator_task(
             catalog (str): task catalog
             inputs (list): task inputs
             outputs (list): task outputs
+            input_model (BaseModel): model declaring the input sockets, checked at every
+                call and again before the body runs
+            output_model (BaseModel): model declaring the output sockets and validating
+                the return value
         """
 
         def decorator(obj: Union[WorkGraph, type, callable]) -> TaskHandle:
             normalized_handlers = normalize_error_handlers(error_handlers)
+            in_spec, out_spec, executor = apply_models(
+                obj, inputs, outputs, input_model, output_model, api=SocketSpecAPI
+            )
+            _refuse_node_typed_inputs(input_model, in_spec)
             spec = _spec_for(
                 obj,
                 identifier=identifier,
                 catalog=catalog,
-                inputs=inputs,
-                outputs=outputs,
+                inputs=in_spec,
+                outputs=out_spec,
                 error_handlers=normalized_handlers,
             )
+            if executor is not obj:
+                # The spec is inferred from the undecorated function, so its
+                # signature, source and return annotation stay visible; only
+                # what runs changes.
+                spec = replace(spec, executor=RuntimeExecutor.from_callable(executor))
 
             handle = TaskHandle(spec)
-            handle._callable = obj
+            handle._callable = executor
             return handle
 
         return decorator
@@ -167,30 +233,39 @@ def decorator_graph(
         outputs: Optional[SocketSpec | list] = None,
         max_depth: int = 100,
         max_number_jobs: Optional[int] = None,
+        input_model: Optional[Type[BaseModel]] = None,
+        output_model: Optional[Type[BaseModel]] = None,
     ) -> Callable:
         """Generate a decorator that register a function as a graph task.
         Attributes:
             indentifier (str): task identifier
             catalog (str): task catalog
             inputs (list): task inputs
             outputs (list): task outputs
+            input_model (BaseModel): model declaring the input sockets, checked at every
+                call and again when the graph is expanded
+            output_model (BaseModel): refused; a graph returns socket references, which
+                stand for values that do not exist yet
         """
 
         def decorator(func) -> TaskHandle:
             from aiida_workgraph.tasks.graph_task import _build_graph_task_taskspec
 
+            in_spec, _, executor = apply_models(
+                func, inputs, None, input_model, output_model, is_graph=True, api=SocketSpecAPI
+            )
             handle = TaskHandle(
                 _build_graph_task_taskspec(
                     func,
                     identifier=identifier,
                     catalog=catalog,
-                    in_spec=inputs,
+                    in_spec=in_spec,
                     out_spec=outputs,
                     max_depth=max_depth,
                     max_number_jobs=max_number_jobs,
                 )
             )
-            handle._callable = func
+            handle._callable = executor
             return handle
 
         return decorator
@@ -202,14 +277,37 @@ def calcfunction(
         outputs: Optional[SocketSpec | list] = None,
         catalog: Optional[str] = None,
         error_handlers: Optional[Dict[str, ErrorHandlerSpec]] = None,
+        input_model: Optional[Type[BaseModel]] = None,
+        output_model: Optional[Type[BaseModel]] = None,
     ) -> Callable:
+        """Generate a decorator registering a function as a calcfunction task.
+
+        Attributes:
+            inputs (list): task inputs
+            outputs (list): task outputs
+            input_model (BaseModel): model declaring the input sockets, checked at every
+                call and again before the body runs; a field declaring an AiiDA type is
+                handed the node, which is what a calcfunction's body receives
+            output_model (BaseModel): model declaring the output sockets and validating
+                the return value
+        """
+
         def decorator(func) -> TaskHandle:
-            func_decorated = calcfunction(func)
+            in_spec, out_spec, executor = apply_models(
+                func, inputs, outputs, input_model, output_model, api=SocketSpecAPI
+            )
+            _refuse_namespace_inputs(input_model, in_spec)
+            # The models are enforced inside the process, so what AiiDA runs is
+            # the wrapper and the nodes it is called with reach the body. The
+            # calcfunction is what the executor has to resolve to, so it takes
+            # over the name the wrapper was bound under.
+            func_decorated = calcfunction(executor)
+            rebind_executor_callable(func_decorated, executor)
             handle = TaskHandle(
                 _build_aiida_function_taskspec(
                     func_decorated,
-                    in_spec=inputs,
-                    out_spec=outputs,
+                    in_spec=in_spec,
+                    out_spec=out_spec,
                     catalog=catalog,
                     error_handlers=error_handlers,
                 )

src/aiida_workgraph/serialization.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum
 from typing import Any, Dict, Optional
 
 from aiida_pythonjob.data.serializer import all_serializers
@@ -8,6 +9,103 @@
 from node_graph.utils import resolve_tagged_values
 
 
+def _flatten_enums(value: Any) -> Any:
+    """Replace ``Enum`` members with their bare ``.value``, recursively.
+
+    ``general_serializer`` (aiida-pythonjob) has no serializer for ``Enum``
+    members: a raw ``Enum`` -- or one nested in a dict/list/tuple -- reaches
+    it and the whole submission fails at submit time with an opaque
+    serialization error. Collapsing members to their ``.value`` here keeps
+    the payload JSON-serializable so ``serialize_ports`` succeeds. The
+    ``isinstance`` check also matches wrapt proxies (node-graph's
+    ``TaggedValue``) wrapping an enum member.
+
+    Flattening is one-way. Whether a body then receives the ``Enum`` it
+    declared is the installed ``node_graph``'s call: the read side
+    (``coerce_inputs_from_spec``) rebuilds a socket's value only from the
+    ``structured_type`` descriptor its spec records, and not every
+    ``node_graph`` records one for ``Enum``. Without it a function task's
+    body receives the bare value and a ``@task.graph`` body the stored
+    ``orm.Str``; with it both receive the member, a ``@task.graph`` body
+    behind a ``TaggedValue``. A body that must work under either writes
+    ``Color(getattr(c, 'value', c))``, never ``Color(c)`` and never
+    ``c is Color.RED``. See
+    ``tests/test_serializer.py::test_enum_arrival_follows_the_node_graph_capability``.
+
+    ``set``/``frozenset`` are deliberately left untouched: a set fails in
+    ``general_serializer`` regardless of its contents (not JSON-serializable,
+    no registered serializer), so descending into one to flatten enums would
+    not make it serializable -- see
+    ``tests/test_serializer.py::test_flatten_leaves_sets_untouched``.
+
+    A container with no Enum anywhere inside it comes back as the exact
+    object passed in, not a rebuilt plain ``dict``/``tuple``. Rebuilding
+    unconditionally would downgrade a ``namedtuple`` to a bare ``tuple`` or
+    an ``OrderedDict``/``defaultdict`` to a plain ``dict`` even when nothing
+    needed flattening -- see
+    ``tests/test_serializer.py::test_flatten_passthrough_preserves_namedtuple_type``
+    and ``::test_flatten_passthrough_preserves_dict_subclass_type``.
+    """
+    if isinstance(value, Enum):
+        return _flatten_enums(value.value)
+    if isinstance(value, dict):
+        out: Dict[Any, Any] = {}
+        changed = False
+        for k, v in value.items():
+            flat_key = _flatten_enums(k)
+            flat_val = _flatten_enums(v)
+            if flat_key is not k or flat_val is not v:
+                changed = True
+            if flat_key in out:
+                # Two distinct keys (e.g. an Enum member and its bare value,
+                # or two members sharing a ``.value``) collapse to one; raise
+                # rather than silently drop an entry.
+                raise ValueError(f'Enum key flattening collision: multiple keys map to {flat_key!r}')
+            out[flat_key] = flat_val
+        return out if changed else value
+    if isinstance(value, list):
+        flat_list = [_flatten_enums(v) for v in value]
+        return flat_list if any(fv is not v for fv, v in zip(flat_list, value)) else value
+    if isinstance(value, tuple):
+        flat_tuple = tuple(_flatten_enums(v) for v in value)
+        return flat_tuple if any(fv is not v for fv, v in zip(flat_tuple, value)) else value
+    return value
+
+
+def _to_declared_python(value: Any) -> Any:
+    """Return ``value`` with the node the write path wrapped it in taken off.
+
+    ``orm.BaseType``, ``orm.Dict`` and ``orm.List`` are the nodes the write
+    path creates for plain Python; every other ``orm.Data`` is a value in its
+    own right and is returned as it is.
+
+    A tagged value is retagged whether or not anything came off it, and keeps
+    its uuid: the tag is what a graph body turns into a link, and the uuid is
+    what makes the body's value and the graph's input one value rather than
+    two. Handing back the wrapped value on the branch where nothing needed
+    unwrapping would drop the tag on exactly the plain fields most bodies
+    take.
+    """
+    from aiida import orm
+    from node_graph.socket import TaggedValue
+
+    if isinstance(value, TaggedValue):
+        tagged = TaggedValue(_to_declared_python(value.__wrapped__), socket=value._socket)
+        tagged._self_uuid = value._uuid
+        return tagged
+    if isinstance(value, orm.BaseType):
+        return value.value
+    if isinstance(value, orm.Dict):
+        return value.get_dict()
+    if isinstance(value, orm.List):
+        return value.get_list()
+    if isinstance(value, dict):
+        return {key: _to_declared_python(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return type(value)(_to_declared_python(item) for item in value)
+    return value
+
+
 class AiidaSerializationAdapter(SerializationAdapter):
     id: str = 'aiida'
     name: str = 'AiiDA'
@@ -17,19 +115,57 @@ def __init__(self, serializers: Optional[Dict[str, str]] = None, user: Any = Non
         self.user = user
 
     def serialize(self, value: Any, socket: Any, *, store: bool) -> Any:
+        from node_graph.input_model import model_dumper_for_socket
+
         if socket is None:
             return value
         spec = socket._to_spec()
         resolve_tagged_values(value)
+        dump = model_dumper_for_socket(socket)
+        if dump is not None:
+            # The task's input model owns this socket's wire form: the model's
+            # own serialization renders the value, so a field type JSON cannot
+            # hold reaches the database through the field_serializer that
+            # declares its form.
+            value = dump(value)
+        else:
+            value = _flatten_enums(value)
         return serialize_ports(
             python_data=value,
             port_schema=spec,
             serializers=self.serializers,
             user=self.user,
         )
 
+    def deserialize(self, value: Any, socket: Any) -> Any:
+        """Give a model-owned socket's value the form the model's field declares.
+
+        The write path promotes a plain Python value to the ``orm`` node that
+        carries it into provenance; this is the read edge that takes the node
+        off again, so a field declared ``str`` reaches the body as ``str``
+        rather than as ``orm.Str``. Which fields those are is the model's call
+        and not the socket identifier's: ``int`` and ``orm.Int`` are the same
+        identifier, and only the model says which of the two was written, so
+        the leaf's ``body_receives`` mark decides. A socket no model owns is
+        left to the base adapter.
+
+        The tag a value wears is put back on, because it is what a graph body
+        turns into a link: unwrapping without it would leave the body holding
+        a copy of the graph's input rather than a reference to it.
+        """
+        from node_graph.input_model import BODY_RECEIVES
+
+        extras = getattr(getattr(socket, '_metadata', None), 'extras', None) or {}
+        arrival = extras.get(BODY_RECEIVES)
+        if arrival is None:
+            return super().deserialize(value, socket)
+        if arrival == 'node':
+            return value
+        return _to_declared_python(value)
+
     def serialize_ports(self, python_data: Any, port_schema: Any, *, store: bool) -> Any:
         resolve_tagged_values(python_data)
+        python_data = _flatten_enums(python_data)
         return serialize_ports(
             python_data=python_data,
             port_schema=port_schema,