diff --git a/src/aiida_workgraph/decorator.py b/src/aiida_workgraph/decorator.py index 7627d553..05f3047c 100644 --- a/src/aiida_workgraph/decorator.py +++ b/src/aiida_workgraph/decorator.py @@ -1,5 +1,7 @@ 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 @@ -7,10 +9,13 @@ 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,6 +233,8 @@ 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: @@ -174,23 +242,30 @@ def decorator_graph( 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, ) diff --git a/src/aiida_workgraph/serialization.py b/src/aiida_workgraph/serialization.py index 92705bdc..834afb34 100644 --- a/src/aiida_workgraph/serialization.py +++ b/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,10 +115,21 @@ 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, @@ -28,8 +137,35 @@ def serialize(self, value: Any, socket: Any, *, store: bool) -> Any: 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, diff --git a/src/aiida_workgraph/socket_spec.py b/src/aiida_workgraph/socket_spec.py index b97cbbce..277ff190 100644 --- a/src/aiida_workgraph/socket_spec.py +++ b/src/aiida_workgraph/socket_spec.py @@ -25,6 +25,8 @@ 'validate_socket_data', 'infer_specs_from_callable', 'from_aiida_process', + 'spec_from_model', + 'node_typed_paths', 'SocketSpecSelect', 'select', 'meta', @@ -148,3 +150,36 @@ def from_aiida_process( validate_socket_data = SocketSpecAPI.validate_socket_data infer_specs_from_callable = SocketSpecAPI.infer_specs_from_callable from_aiida_process = SocketSpecAPI.from_aiida_process + + +def spec_from_model(model: Any) -> SocketSpec: + """Return the socket namespace ``model`` describes, in this package's vocabulary. + + ``node_graph.input_model.spec_from_model`` builds the same namespace out of + node-graph's socket identifiers; this binds it to aiida-workgraph's. + """ + from node_graph.input_model import spec_from_model as _spec_from_model + + return _spec_from_model(model, SocketSpecAPI) + + +def node_typed_paths(spec: SocketSpec, prefix: Tuple[str, ...] = ()) -> list[str]: + """Return the path of every leaf of ``spec`` a model says arrives as a node. + + The walk crosses namespaces and the item shape of a dynamic namespace, so + a field declared deep inside a nested model, or inside the items of a + ``dict[str, T]``, is found where it sits. + """ + from node_graph.input_model import BODY_RECEIVES + + if spec is None: + return [] + if not spec.is_namespace(): + extras = getattr(spec.meta, 'extras', None) or {} + return ['.'.join(prefix)] if extras.get(BODY_RECEIVES) == 'node' else [] + found: list[str] = [] + for name, child in (spec.fields or {}).items(): + found.extend(node_typed_paths(child, prefix + (name,))) + if spec.item is not None: + found.extend(node_typed_paths(spec.item, prefix + ('',))) + return found diff --git a/tests/test_input_model.py b/tests/test_input_model.py new file mode 100644 index 00000000..670abcc7 --- /dev/null +++ b/tests/test_input_model.py @@ -0,0 +1,1258 @@ +"""A Pydantic model as a task's wire contract, under the AiiDA engine. + +The contract itself -- which sockets a model declares, and when its rules are +held to -- is node-graph's, and its tests live there. What is asserted here is +what only this package can answer: that the model decides the *stored* form of +a value, that the body gets the rich object back out of storage, and that a +rule broken inside a submitted process fails that process with a message +naming the task. +""" + +from __future__ import annotations + +import enum +from decimal import Decimal +from typing import Any + +import pytest +from aiida import orm +from node_graph.input_model import BODY_RECEIVES, ModelContractError, TaskInputValidationError +from pydantic import ( + BaseModel, + ConfigDict, + Field, + field_serializer, + field_validator, + model_validator, +) + +from aiida_workgraph import WorkGraph, task +from aiida_workgraph.socket_spec import spec_from_model + +#: aiida-pythonjob's exit status for a body that raised. +FUNCTION_FAILED = 323 + + +class Color(enum.Enum): + RED = 'red' + BLUE = 'blue' + + +# -------------------------------------------------------------------------- +# 1. The model builds this package's sockets +# -------------------------------------------------------------------------- + + +class AddInputs(BaseModel): + """Two summands, the second optional.""" + + x: int + y: int = 7 + + +@task(input_model=AddInputs) +def add(x, y): + return x + y + + +def test_the_sockets_carry_this_packages_identifiers(): + """The spec is built through aiida-workgraph's socket vocabulary, not node-graph's.""" + fields = add._spec.inputs.fields + assert fields['x'].identifier == 'workgraph.int' + assert fields['x'].meta.required is True + assert fields['y'].identifier == 'workgraph.int' + assert fields['y'].meta.required is False + assert fields['y'].default == 7 + + +def test_an_omitted_input_runs_on_the_models_default(): + wg = WorkGraph('add_default') + node = wg.add_task(add, name='add', x=2) + wg.run() + assert node.outputs.result.value.value == 9 + + +class NudgeInputs(BaseModel): + """``by`` defaults to ``None``, the value the engine declines to store.""" + + x: int + by: int | None = None + + +@task(input_model=NudgeInputs) +def nudge(x, by): + return x if by is None else x + by + + +def test_a_field_defaulting_to_none_is_not_a_missing_required_input(): + assert nudge._spec.inputs.fields['by'].meta.required is False + wg = WorkGraph('nudge') + node = wg.add_task(nudge, name='nudge', x=4) + wg.run() + assert node.outputs.result.value.value == 4 + + +def test_a_model_on_a_process_task_is_refused(): + from aiida.calculations.arithmetic.add import ArithmeticAddCalculation + + with pytest.raises(ModelContractError, match='plain Python function task'): + task(input_model=AddInputs)(ArithmeticAddCalculation) + + +# -------------------------------------------------------------------------- +# 2. An Enum: stored as its bare value, delivered as the member +# -------------------------------------------------------------------------- + + +class PaintInputs(BaseModel): + color: Color + + +@task(input_model=PaintInputs, outputs=['is_member', 'seen']) +def paint(color): + # ``is`` and not ``==``: the body gets the member itself, not a look-alike. + return {'is_member': color is Color.RED, 'seen': type(color).__name__} + + +@task(inputs=spec_from_model(PaintInputs), outputs=['is_member', 'seen']) +def paint_without_model(color): + """Same sockets, no contract: the control for :func:`paint`.""" + return {'is_member': color is Color.RED, 'seen': type(color).__name__} + + +def test_an_enum_is_stored_as_its_bare_value(): + wg = WorkGraph('paint_store') + node = wg.add_task(paint, name='paint', color=Color.RED) + wg.run() + stored = dict(node.process.inputs.function_inputs)['color'] + assert stored.value == 'red' + + +def test_the_body_receives_the_enum_member_and_only_with_the_model(): + wg = WorkGraph('paint_body') + with_model = wg.add_task(paint, name='with_model', color=Color.RED) + without_model = wg.add_task(paint_without_model, name='without_model', color=Color.RED) + wg.run() + + assert with_model.outputs.is_member.value.value is True + assert with_model.outputs.seen.value.value == 'Color' + # The control proves the model is what rebuilt the member: the same socket + # hands the bare value to a body no model stands in front of. + assert without_model.outputs.is_member.value.value is False + assert without_model.outputs.seen.value.value == 'str' + + +def test_an_enum_survives_a_workgraph_round_trip(): + wg = WorkGraph('paint_round_trip') + wg.add_task(paint, name='paint', color=Color.BLUE) + rebuilt = WorkGraph.from_dict(wg.to_dict()) + assert rebuilt.tasks.paint.inputs.color.value == Color.BLUE + rebuilt.run() + assert rebuilt.tasks.paint.outputs.is_member.value.value is False + assert rebuilt.tasks.paint.outputs.seen.value.value == 'Color' + + +def test_no_class_path_is_written_into_the_task(): + """The spec stored with the task names no class: the model is reached through the code.""" + assert 'structured_type' not in paint._spec.inputs.fields['color'].meta.extras + + +# -------------------------------------------------------------------------- +# 3. A type JSON cannot hold, carried by the model's own serializer +# -------------------------------------------------------------------------- + + +class MoneyInputs(BaseModel): + """A ``Decimal`` amount, stored as the string the model renders.""" + + amount: Decimal + + @field_serializer('amount') + def _dump_amount(self, value: Decimal) -> str: + return str(value) + + @field_validator('amount', mode='before') + @classmethod + def _load_amount(cls, value): + return value if isinstance(value, Decimal) else Decimal(str(value)) + + +@task(input_model=MoneyInputs, outputs=['kind', 'doubled']) +def double_money(amount): + return {'kind': type(amount).__name__, 'doubled': str(amount * 2)} + + +@task(inputs=spec_from_model(MoneyInputs), outputs=['kind', 'doubled']) +def double_money_without_model(amount): + """Same sockets, no contract: the control for :func:`double_money`.""" + return {'kind': type(amount).__name__, 'doubled': str(amount * 2)} + + +def test_a_field_serializer_decides_the_stored_form(): + wg = WorkGraph('money_store') + node = wg.add_task(double_money, name='money', amount=Decimal('0.10')) + wg.run() + stored = dict(node.process.inputs.function_inputs)['amount'] + assert stored.value == '0.10' + + +def test_the_body_receives_the_decimal_and_only_with_the_model(): + wg = WorkGraph('money_body') + node = wg.add_task(double_money, name='money', amount=Decimal('0.10')) + wg.run() + assert node.outputs.kind.value.value == 'Decimal' + # Exact, because a Decimal round-tripped as a string never became a float. + assert node.outputs.doubled.value.value == '0.20' + + +def test_without_the_model_the_same_value_cannot_even_be_stored(): + """The control: nothing else in the stack knows how to write a ``Decimal``.""" + wg = WorkGraph('money_control') + wg.add_task(double_money_without_model, name='money', amount=Decimal('0.10')) + with pytest.raises(ValueError, match='decimal.Decimal'): + wg.run() + + +# -------------------------------------------------------------------------- +# 4. Rules the socket layer cannot see fail the process that broke them +# -------------------------------------------------------------------------- + + +class RangeInputs(BaseModel): + """``low`` and ``high`` are ints the socket accepts; their order is the model's rule.""" + + low: int + high: int = Field(le=100) + + @model_validator(mode='after') + def _ordered(self): + if self.low >= self.high: + raise ValueError('low must be below high') + return self + + +@task(input_model=RangeInputs) +def span(low, high): + return high - low + + +def test_a_cross_field_rule_fails_the_process_naming_the_task(): + wg = WorkGraph('span_bad_order') + node = wg.add_task(span, name='span', low=9, high=3) + wg.run() + assert node.process.exit_status == FUNCTION_FAILED + assert "Task 'span' got inputs RangeInputs rejects" in node.process.exit_message + assert 'low must be below high' in node.process.exit_message + + +def test_a_field_constraint_fails_where_it_is_written(): + """``le=100`` is a model rule the socket layer cannot see, and ``add_task`` still refuses 500.""" + wg = WorkGraph('span_too_high') + with pytest.raises(TaskInputValidationError, match='less than or equal to 100'): + wg.add_task(span, name='span', low=1, high=500) + + +def test_without_checkpoint_a_that_constraint_reaches_the_run_edge(monkeypatch): + """The control: with the write unchecked, ``le=100`` is first seen where the body runs.""" + from node_graph import input_model + + monkeypatch.setattr(input_model, 'validate_task_inputs', lambda task, inputs: None) + wg = WorkGraph('span_too_high_unchecked') + node = wg.add_task(span, name='span', low=1, high=500) + wg.run() + assert node.process.exit_status == FUNCTION_FAILED + assert 'less than or equal to 100' in node.process.exit_message + + +# -------------------------------------------------------------------------- +# 5. A graph task's contract, held where the engine expands it +# -------------------------------------------------------------------------- + + +class WindowInputs(BaseModel): + """The window a graph opens; ``lower`` below ``upper`` is the graph's own rule.""" + + lower: int + upper: int + + @model_validator(mode='after') + def _ordered(self): + if self.lower >= self.upper: + raise ValueError('lower must be below upper') + return self + + +@task.graph(input_model=WindowInputs) +def window(lower, upper): + return add(x=lower, y=upper) + + +@task() +def shrink(value): + return value - 10 + + +@task.graph() +def window_of_a_computed_bound(lower, upper): + """The subgraph's bound is decided by a task, so it is a value only at run time.""" + shrunk = shrink(value=upper) + return window(lower=lower, upper=shrunk.result) + + +def test_a_graph_contract_holds_when_the_graph_is_submitted(): + wg = WorkGraph('window_ok') + node = wg.add_task(window, name='window', lower=1, upper=3) + wg.run() + assert node.process.exit_status == 0 + + +def test_a_graph_contract_fails_the_submitted_graph(): + """The graph is refused as it is expanded, so it never becomes a process at all.""" + wg = WorkGraph('window_bad') + node = wg.add_task(window, name='window', lower=9, upper=3) + wg.run() + assert node.process is None + assert node.state == 'FAILED' + assert wg.process.exit_status != 0 + + +class Named(BaseModel): + """A `str` field, which is what the engine's own wrappers get in the way of.""" + + label: str + count: int + + +@task() +def echo(label, count): + return f'{label}-{count}' + + +@task.graph(input_model=Named) +def named(label, count): + return echo(label=label, count=count) + + +def test_a_graph_contract_reads_through_the_engines_wrappers(): + """A graph body is handed storage nodes; the contract is checked against what they hold.""" + wg = WorkGraph('named_ok') + node = wg.add_task(named, name='named', label='silicon', count=2) + wg.run() + assert node.process.exit_status == 0 + + +class Priced(BaseModel): + """A `Decimal` beside a `str`: one kind no socket identifier can carry.""" + + label: str + amount: Decimal + + +@task() +def show(label, amount): + return f'{label}-{amount}' + + +@task.graph(input_model=Priced) +def priced(label, amount): + return show(label=label, amount=str(amount)) + + +def test_a_graph_contract_reads_a_field_no_identifier_carries(): + """`Decimal` is stored as the string the model rendered, and read back through the model.""" + wg = WorkGraph('priced_ok') + node = wg.add_task(priced, name='priced', label='silicon', amount=Decimal('1.50')) + wg.run() + assert node.process.exit_status == 0 + + +def test_without_the_unwrap_that_field_would_be_refused(monkeypatch): + """The control: leave the node on and `Decimal` refuses the `orm.Str` holding its rendering. + + The field is deliberately one a socket identifier cannot answer for -- it is + `workgraph.annotated`, not `workgraph.string` -- so the control still + discriminates under an adapter that unwraps by identifier. + """ + from node_graph.serializer import SerializationAdapter + + from aiida_workgraph.serialization import AiidaSerializationAdapter + + monkeypatch.setattr(AiidaSerializationAdapter, 'deserialize', SerializationAdapter.deserialize) + wg = WorkGraph('priced_unread') + node = wg.add_task(priced, name='priced', label='silicon', amount=Decimal('1.50')) + wg.run() + assert node.process is None + assert node.state == 'FAILED' + + +ARRIVED: dict = {} + + +def _arrival(value): + """Return the type of what a body was handed, seeing through the tag it wears.""" + return type(getattr(value, '__wrapped__', value)).__name__ + + +class EveryKind(BaseModel): + """One field per kind the read edge has to tell apart.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + text: str + number: int + fraction: float + flag: bool + mapping: dict + items: list + anything: Any + node: orm.Int + + +@task.graph(input_model=EveryKind) +def records_arrivals(text, number, fraction, flag, mapping, items, anything, node): + ARRIVED.clear() + ARRIVED.update( + text=_arrival(text), + number=_arrival(number), + fraction=_arrival(fraction), + flag=_arrival(flag), + mapping=_arrival(mapping), + items=_arrival(items), + anything=_arrival(anything), + node=_arrival(node), + ) + + +def test_a_body_receives_what_its_field_declares(): + """Declared Python arrives as Python; `Any` and an AiiDA type arrive as nodes.""" + wg = WorkGraph('arrivals') + wg.add_task( + records_arrivals, + name='r', + text='silicon', + number=3, + fraction=1.5, + flag=True, + mapping={'k': 1}, + items=[1, 2], + anything='whatever', + node=orm.Int(7), + ) + wg.run() + assert ARRIVED == { + 'text': 'str', + 'number': 'int', + 'fraction': 'float', + 'flag': 'bool', + 'mapping': 'dict', + 'items': 'list', + 'anything': 'Str', + 'node': 'Int', + } + + +def test_the_socket_identifier_alone_cannot_tell_int_from_orm_int(): + """Both are ``workgraph.int``; only the model says which of the two was declared.""" + from aiida_workgraph.socket_spec import spec_from_model + + spec = spec_from_model(EveryKind) + assert spec.fields['number'].identifier == spec.fields['node'].identifier + assert spec.fields['number'].meta.extras[BODY_RECEIVES] == 'python' + assert spec.fields['node'].meta.extras[BODY_RECEIVES] == 'node' + + +def _naive_unwrap(value): + """Unwrap the way an identifier-keyed read edge does, handing the tag back to no one.""" + from node_graph.socket import TaggedValue + + plain = value.__wrapped__ if isinstance(value, TaggedValue) else value + if isinstance(plain, orm.BaseType): + return plain.value + if isinstance(plain, orm.Dict): + return plain.get_dict() + if isinstance(plain, orm.List): + return plain.get_list() + return plain + + +def _label_nodes(node): + """Return the node written into the graph's ``label`` and the one its subtask got.""" + outer = inner = None + for link in node.process.base.links.get_incoming().all(): + if link.link_label.endswith('label'): + outer = link.node + for child in node.process.called: + for link in child.base.links.get_incoming().all(): + if link.link_label.endswith('label'): + inner = link.node + return outer, inner + + +def test_an_unwrapped_value_keeps_the_tag_that_draws_its_link(): + """The body forwards a `str` field, and the child must be linked to it, not given a copy. + + Provenance is what says which: the node the subtask reads has to be the + very node the graph was given, not a second one holding the same string. + """ + wg = WorkGraph('linked') + node = wg.add_task(named, name='named', label='silicon', count=2) + wg.run() + assert node.process.exit_status == 0 + outer, inner = _label_nodes(node) + assert outer is not None and inner is not None + assert outer.uuid == inner.uuid + + +def test_without_the_tag_the_subtask_reads_a_node_nobody_produced(monkeypatch): + """The control: unwrap without retagging and the body holds a copy, not a reference.""" + import aiida_workgraph.serialization as serialization + + monkeypatch.setattr(serialization, '_to_declared_python', _naive_unwrap) + wg = WorkGraph('linked_naive') + node = wg.add_task(named, name='named', label='silicon', count=2) + wg.run() + outer, inner = _label_nodes(node) + assert outer is not None and inner is not None + assert outer.uuid != inner.uuid + assert len(inner.base.links.get_incoming().all()) == 0 + + +def test_a_value_that_needed_no_unwrapping_keeps_its_tag_too(): + """A plain `str` has no node to take off, and dropping the tag there drops the link.""" + from node_graph.socket import TaggedValue + + from aiida_workgraph.serialization import _to_declared_python + + tagged = TaggedValue('silicon', socket=object()) + unwrapped = _to_declared_python(tagged) + assert isinstance(unwrapped, TaggedValue) + assert unwrapped._socket is tagged._socket + assert unwrapped._uuid == tagged._uuid + + +def test_a_value_that_was_unwrapped_keeps_the_uuid_it_arrived_with(): + """One value, one uuid: a new one would make the body's value a second value.""" + from node_graph.socket import TaggedValue + + from aiida_workgraph.serialization import _to_declared_python + + tagged = TaggedValue(orm.Int(7), socket=object()) + unwrapped = _to_declared_python(tagged) + assert unwrapped == 7 + assert unwrapped._uuid == tagged._uuid + + +def test_a_runtime_value_is_checked_at_the_graph_it_reaches(): + """Nothing knows ``upper`` is 5 until ``shrink`` has run, so this is the first chance.""" + wg = WorkGraph('window_computed') + node = wg.add_task(window_of_a_computed_bound, name='outer', lower=9, upper=15) + wg.run() + assert node.process.exit_status != 0 + + +class Spin(enum.Enum): + """The four spin treatments a workflow accepts.""" + + NONE = 'none' + COLLINEAR = 'collinear' + NON_COLLINEAR = 'non_collinear' + SPIN_ORBIT = 'spin_orbit' + + +class DielectricInputs(BaseModel): + """Two of the four members are ph.x's own limit, written where ph.x is declared.""" + + spin: Spin = Spin.NONE + structure: str + + @field_validator('spin') + @classmethod + def _supported(cls, value): + if value in (Spin.NON_COLLINEAR, Spin.SPIN_ORBIT): + raise ValueError('ph.x has no electric-field perturbation for noncollinear magnetism') + return value + + +@task(input_model=DielectricInputs, outputs=['seen']) +def dielectric(spin, structure): + return {'seen': type(spin).__name__} + + +class EverySpin(BaseModel): + spin: Spin = Spin.NONE + structure: str + + +@task.graph(input_model=EverySpin) +def eps(spin, structure): + return dielectric(spin=spin, structure=structure).seen + + +def test_a_rule_on_an_inner_task_fails_the_expansion(): + """The graph takes every spin, so the value meets ph.x's rule as the body wires it.""" + wg = WorkGraph('eps_noncollinear') + node = wg.add_task(eps, name='eps', spin=Spin.NON_COLLINEAR, structure='si') + wg.run() + assert node.state == 'FAILED' + assert node.process is None + # Neither the subgraph nor the task it would have held became a process. + assert [called.process_label for called in wg.process.called_descendants] == [] + + +def test_the_spin_the_rule_admits_reaches_the_body_as_the_member(): + """The control: the same wiring expands, runs, and the enum survives storage.""" + wg = WorkGraph('eps_collinear') + node = wg.add_task(eps, name='eps', spin=Spin.COLLINEAR, structure='si') + wg.run() + assert node.process.exit_status == 0 + assert sorted(called.process_label for called in wg.process.called_descendants) == [ + 'WorkGraph', + 'dielectric', + ] + assert node.outputs.result.value.value == 'Spin' + + +# -------------------------------------------------------------------------- +# 6. Output models +# -------------------------------------------------------------------------- + + +class SumAndProduct(BaseModel): + sum: int + product: int + + +@task(output_model=SumAndProduct) +def combine(x, y): + return {'sum': x + y, 'product': x * y} + + +@task(output_model=SumAndProduct) +def combine_forgetting_product(x, y): + return {'sum': x + y} + + +@task(output_model=SumAndProduct) +def combine_with_a_bad_type(x, y): + return {'sum': 'not a number', 'product': x * y} + + +def test_the_output_sockets_come_from_the_output_model(): + assert set(combine._spec.outputs.fields) == {'sum', 'product'} + + +def test_a_return_the_model_accepts_lands_on_the_sockets(): + wg = WorkGraph('combine_ok') + node = wg.add_task(combine, name='combine', x=2, y=3) + wg.run() + assert node.outputs.sum.value.value == 5 + assert node.outputs.product.value.value == 6 + + +def test_a_missing_output_fails_at_the_source_task(): + wg = WorkGraph('combine_missing') + node = wg.add_task(combine_forgetting_product, name='combine', x=2, y=3) + wg.run() + assert node.process.exit_status == FUNCTION_FAILED + assert "Task 'combine_forgetting_product' returned outputs SumAndProduct rejects" in node.process.exit_message + assert 'product' in node.process.exit_message + + +def test_a_mistyped_output_fails_at_the_source_task(): + wg = WorkGraph('combine_bad_type') + node = wg.add_task(combine_with_a_bad_type, name='combine', x=2, y=3) + wg.run() + assert node.process.exit_status == FUNCTION_FAILED + assert 'sum' in node.process.exit_message + + +# -------------------------------------------------------------------------- +# 7. Mappings whose size is only known at runtime +# -------------------------------------------------------------------------- + + +class Block(BaseModel): + """One member of a mapping the task fills in at runtime.""" + + width: int + label: str + + +class Blocks(BaseModel): + blocks: dict[str, Block] + + +class Recipe(BaseModel): + """Two mappings on the way in: one of plain strings, one needing the model.""" + + pseudos: dict[str, str] + amounts: dict[str, Decimal] = {} + + @field_serializer('amounts') + def _dump_amounts(self, value: dict[str, Decimal]) -> dict[str, str]: + return {key: str(amount) for key, amount in value.items()} + + @field_validator('amounts', mode='before') + @classmethod + def _load_amounts(cls, value): + return {key: amount if isinstance(amount, Decimal) else Decimal(str(amount)) for key, amount in value.items()} + + +@task(output_model=Blocks) +def cut_blocks(n): + """Decide how many members the mapping has, inside the body.""" + return {'blocks': {f'b{i}': {'width': i + 1, 'label': f'block {i}'} for i in range(n)}} + + +@task(output_model=Blocks) +def cut_one_bad_block(n): + return {'blocks': {'b0': {'width': 1, 'label': 'fine'}, 'b1': {'width': 'wide', 'label': 'bad'}}} + + +@task(input_model=Blocks, outputs=['widest']) +def widest_block(blocks): + # Each member arrived through ``Block`` as the keys that were written. + return {'widest': max(item['width'] for item in blocks.values())} + + +@task(input_model=Recipe, outputs=['names', 'kinds']) +def describe_recipe(pseudos, amounts): + return { + 'names': ','.join(sorted(pseudos)), + 'kinds': ','.join(sorted({type(amount).__name__ for amount in amounts.values()})), + } + + +def test_a_typed_mapping_field_becomes_a_dynamic_namespace(): + blocks = cut_blocks._spec.outputs.fields['blocks'] + assert blocks.is_namespace() + assert blocks.meta.dynamic is True + # Every member is typed, so a key's sockets are known before any key is. + assert set(blocks.item.fields) == {'width', 'label'} + assert blocks.item.fields['width'].identifier == 'workgraph.int' + + +def test_a_mapping_becomes_one_socket_per_key_after_the_run(): + wg = WorkGraph('blocks_out') + node = wg.add_task(cut_blocks, name='cut', n=3) + wg.run() + assert set(node.outputs.blocks._sockets) == {'b0', 'b1', 'b2'} + assert node.outputs.blocks.b1.width.value.value == 2 + assert node.outputs.blocks.b1.label.value.value == 'block 1' + + +def test_a_bad_member_fails_at_the_source_task_naming_its_key(): + wg = WorkGraph('blocks_bad') + node = wg.add_task(cut_one_bad_block, name='cut', n=2) + wg.run() + assert node.process.exit_status == FUNCTION_FAILED + # The key is in the path pydantic reports, so the reader knows which member. + assert 'blocks.b1.width' in node.process.exit_message + + +def test_a_downstream_task_consumes_the_mapping(): + wg = WorkGraph('blocks_chain') + source = wg.add_task(cut_blocks, name='cut', n=3) + consumer = wg.add_task(widest_block, name='widest', blocks=source.outputs.blocks) + wg.run() + assert consumer.outputs.widest.value.value == 3 + + +def test_addressing_one_future_member_at_build_time_is_not_available_yet(): + """The per-key sockets appear when the task runs, so a build-time link cannot name one. + + Ordered and by-name member access on a socket that has not produced its + members yet is scinode/node-graph#160. + """ + wg = WorkGraph('blocks_future') + source = wg.add_task(cut_blocks, name='cut', n=3) + with pytest.raises(AttributeError, match="has no sub-socket 'b1'"): + source.outputs.blocks.b1 + + +def test_a_mapping_input_becomes_one_socket_per_key(): + wg = WorkGraph('recipe_in') + node = wg.add_task(describe_recipe, name='recipe', pseudos={'Si': 'si.upf', 'O': 'o.upf'}) + assert set(node.inputs.pseudos._sockets) == {'Si', 'O'} + wg.run() + assert node.outputs.names.value.value == 'O,Si' + + +def test_a_member_of_a_mapping_is_stored_through_the_model(): + wg = WorkGraph('recipe_decimal') + node = wg.add_task( + describe_recipe, + name='recipe', + pseudos={'Si': 'si.upf'}, + amounts={'Si': Decimal('1.50'), 'O': Decimal('0.25')}, + ) + wg.run() + stored = dict(node.process.inputs.function_inputs)['amounts'] + assert {key: value.value for key, value in dict(stored).items()} == {'Si': '1.50', 'O': '0.25'} + assert node.outputs.kinds.value.value == 'Decimal' + + +class Amount(BaseModel): + """A model used as a field, carrying a type JSON cannot hold.""" + + value: Decimal + + @field_serializer('value') + def _dump_value(self, value: Decimal) -> str: + return str(value) + + @field_validator('value', mode='before') + @classmethod + def _load_value(cls, value): + return value if isinstance(value, Decimal) else Decimal(str(value)) + + +class Nested(BaseModel): + cfg: Amount + + +class Mapped(BaseModel): + cfgs: dict[str, Amount] + + +class Layer(BaseModel): + inner: Amount + + +class DeepMapped(BaseModel): + items: dict[str, Layer] + + +@task(input_model=Nested, outputs=['kind', 'stored']) +def read_nested(cfg): + # A nested model reaches the body as the members that were written. + return {'kind': type(cfg['value']).__name__, 'stored': str(cfg['value'])} + + +@task(input_model=Mapped, outputs=['kind', 'stored']) +def read_mapped(cfgs): + # Each item of a mapping reaches the body as the members written into it. + return {'kind': type(cfgs['a']['value']).__name__, 'stored': str(cfgs['a']['value'])} + + +@task(input_model=DeepMapped, outputs=['kind', 'stored']) +def read_deep(items): + value = items['a']['inner']['value'] + return {'kind': type(value).__name__, 'stored': str(value)} + + +@pytest.mark.parametrize( + 'entry_point, payload, path', + [ + (read_nested, {'cfg': {'value': Decimal('1.50')}}, ('cfg', 'value')), + (read_mapped, {'cfgs': {'a': {'value': Decimal('1.50')}}}, ('cfgs', 'a', 'value')), + (read_deep, {'items': {'a': {'inner': {'value': Decimal('1.50')}}}}, ('items', 'a', 'inner', 'value')), + ], + ids=['nested-model', 'mapping-of-models', 'mapping-of-nested-models'], +) +def test_a_model_renders_its_own_leaf_however_deep_it_sits(entry_point, payload, path): + """Whatever models and mappings the path crosses, the model declaring the leaf renders it.""" + wg = WorkGraph(f'depth_{path[0]}') + node = wg.add_task(entry_point, name='read', **payload) + wg.run() + + assert node.process.exit_status == 0 + stored = node.process.inputs.function_inputs + for name in path: + stored = stored[name] + assert stored.value == '1.50' + # And the body gets the Decimal back, not the string that was stored. + assert node.outputs.kind.value.value == 'Decimal' + assert node.outputs.stored.value.value == '1.50' + + +# -------------------------------------------------------------------------- +# 8. What a task without a model does +# -------------------------------------------------------------------------- + + +@task +def plain_add(x: int, y: int = 7) -> int: + return x + y + + +def test_a_task_without_a_model_is_untouched(): + fields = plain_add._spec.inputs.fields + assert fields['x'].identifier == 'workgraph.int' + assert fields['y'].default == 7 + wg = WorkGraph('plain') + node = wg.add_task(plain_add, name='add', x=2) + wg.run() + assert node.outputs.result.value.value == 9 + + +# -------------------------------------------------------------------------- +# 9. The node a body was promised, and the edge that takes it off +# -------------------------------------------------------------------------- + + +class StructureInputs(BaseModel): + """One socket the model can only be satisfied by the node itself.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + structure: orm.StructureData + + +class PayloadInputs(BaseModel): + """``Any`` declares nothing to rebuild, so its socket arrives as stored.""" + + payload: Any = None + + +@task.calcfunction(input_model=StructureInputs) +def reads_a_structure(structure): + return orm.Str(type(structure).__name__) + + +@task.calcfunction(input_model=PayloadInputs) +def reads_a_payload(payload): + return orm.Str(type(payload).__name__) + + +@task(outputs=['seen']) +def reads_a_structure_by_annotation(structure: orm.StructureData): + return {'seen': type(structure).__name__} + + +def a_silicon_structure() -> orm.StructureData: + """Return a two-atom silicon cell, stored.""" + structure = orm.StructureData(cell=[[0.0, 2.7, 2.7], [2.7, 0.0, 2.7], [2.7, 2.7, 0.0]]) + structure.append_atom(position=(0.0, 0.0, 0.0), symbols='Si') + structure.append_atom(position=(1.35, 1.35, 1.35), symbols='Si') + return structure.store() + + +def test_the_spec_says_a_node_typed_socket_reaches_the_body_as_the_node(): + """What the contract promises for these two sockets, before anything runs.""" + structure_spec = spec_from_model(StructureInputs) + payload_spec = spec_from_model(PayloadInputs) + assert structure_spec.fields['structure'].meta.extras[BODY_RECEIVES] == 'node' + assert payload_spec.fields['payload'].meta.extras[BODY_RECEIVES] == 'node' + + +def test_a_node_typed_field_is_refused_on_a_pyfunction(): + """A PyFunction reads its inputs out of their nodes, so this field cannot be met.""" + with pytest.raises(ModelContractError, match='@task.calcfunction'): + + @task(input_model=StructureInputs, outputs=['seen']) + def reads_it(structure): + return {'seen': type(structure).__name__} + + +def test_a_field_under_any_is_refused_on_a_pyfunction_too(): + """``Any`` declares nothing to rebuild, so its socket is node-typed as well.""" + with pytest.raises(ModelContractError, match='@task.calcfunction'): + + @task(input_model=PayloadInputs, outputs=['seen']) + def reads_it(payload): + return {'seen': type(payload).__name__} + + +def test_a_calcfunction_body_is_handed_the_node_the_model_declared(): + """Where the rule sends such a field: the body reads the node itself.""" + wg = WorkGraph('structure_leaf') + node = wg.add_task(reads_a_structure, name='leaf', structure=a_silicon_structure()) + wg.run() + assert node.process.exit_status == 0 + assert node.outputs.result.value.value == 'StructureData' + + +def test_a_calcfunction_is_handed_the_node_under_any_as_well(): + """The field kind no per-type override could cover reaches the body whole.""" + wg = WorkGraph('payload_leaf') + node = wg.add_task(reads_a_payload, name='leaf', payload=orm.Dict(dict={'a': 1}).store()) + wg.run() + assert node.process.exit_status == 0 + assert node.outputs.result.value.value == 'Dict' + + +def test_the_same_socket_declared_without_a_model_loses_the_node_in_silence(): + """The control: the edge behaves the same way with no model to notice it.""" + wg = WorkGraph('structure_annotated') + node = wg.add_task(reads_a_structure_by_annotation, name='leaf', structure=a_silicon_structure()) + wg.run() + assert node.process.exit_status == 0 + assert node.outputs.seen.value.value == 'Atoms' + + +# -------------------------------------------------------------------------- +# 10. A graph body is handed the type its field declares +# -------------------------------------------------------------------------- + + +class SpinOnly(BaseModel): + spin: Spin = Spin.NONE + + +@task(outputs=['seen']) +def report(text): + return {'seen': text} + + +def describe(value) -> str: + """Return what a body can tell about the value it was handed.""" + from node_graph.socket import TaggedValue + + # By kind, not by class name: which class tags a value is the tagging + # layer's business, and it has more than one. + kind = 'tagged' if isinstance(value, TaggedValue) else type(value).__name__ + return f'{kind}|{value.__class__.__name__}|{value == Spin.COLLINEAR}|{value in (Spin.NONE, Spin.COLLINEAR)}' + + +@task.graph(input_model=SpinOnly) +def modelled_spin(spin): + return report(text=describe(spin)).seen + + +@task.graph() +def annotated_spin(spin: Spin = Spin.NONE): + return report(text=describe(spin)).seen + + +def test_a_modelled_graph_body_is_handed_the_member_under_its_tag(): + """The member, so ``==`` and ``in`` answer; the tag, so the link is still drawn.""" + wg = WorkGraph('spin_modelled') + node = wg.add_task(modelled_spin, name='g', spin=Spin.COLLINEAR) + wg.run() + assert node.outputs.result.value.value == 'tagged|Spin|True|True' + + +def test_a_graph_body_declaring_the_same_field_by_annotation_agrees(): + """The control: the two ways of declaring one field hand the body one value.""" + wg = WorkGraph('spin_annotated') + node = wg.add_task(annotated_spin, name='g', spin=Spin.COLLINEAR) + wg.run() + assert node.outputs.result.value.value == 'tagged|Spin|True|True' + + +# -------------------------------------------------------------------------- +# 11. A handle that does not take the name of the function it decorates +# -------------------------------------------------------------------------- + + +class BandWindow(BaseModel): + """A field rule for the write, and a cross-field rule only the run edge holds.""" + + lower: int = 0 + upper: int = 1 + + @field_validator('lower') + @classmethod + def _counted(cls, value): + if value < 0: + raise ValueError('lower counts bands, so it cannot be negative') + return value + + @model_validator(mode='after') + def _ordered(self): + if self.upper <= self.lower: + raise ValueError('upper must be above lower') + return self + + +def _bands_body(lower, upper): + return upper - lower + + +#: The decorated name is never rebound, so the module still binds the function. +named_apart_bands = task(input_model=BandWindow)(_bands_body) + + +@task(input_model=BandWindow) +def rebound_bands(lower, upper): + """The ordinary spelling: the handle replaces the module global.""" + return upper - lower + + +SPELLINGS = pytest.mark.parametrize( + 'handle', + [named_apart_bands, rebound_bands], + ids=['handle named apart', 'handle rebound to the name'], +) + + +@SPELLINGS +def test_the_write_is_refused_whatever_the_handle_is_called(handle): + """The write reads the model off the executor, and both spellings store one.""" + wg = WorkGraph('bands_write') + with pytest.raises(TaskInputValidationError, match='cannot be negative'): + wg.add_task(handle, name='w', lower=-5, upper=9) + + +@SPELLINGS +def test_the_submitted_task_holds_the_rule_the_write_could_not(handle): + """A cross-field rule waits for the run edge, which the engine reaches the same way.""" + wg = WorkGraph('bands_run') + node = wg.add_task(handle, name='w', lower=9, upper=3) + wg.run() + assert node.process.exit_status == FUNCTION_FAILED + assert 'upper must be above lower' in node.process.exit_message + + +@SPELLINGS +def test_the_bounds_the_model_admits_still_run(handle): + """The control: what the model accepts reaches the body and comes back.""" + wg = WorkGraph('bands_ok') + node = wg.add_task(handle, name='w', lower=1, upper=9) + wg.run() + assert node.process.exit_status == 0 + assert node.outputs.result.value.value == 8 + + +def test_the_process_is_still_labelled_with_the_name_of_its_function(): + """What the executor is stored under is its own business, not the label's.""" + wg = WorkGraph('bands_label') + node = wg.add_task(named_apart_bands, name='w', lower=1, upper=9) + wg.run() + assert node.process.process_label == '_bands_body' + + +# -------------------------------------------------------------------------- +# 12. What a member nobody wrote is worth at this engine's run edge +# -------------------------------------------------------------------------- + + +class EngineSystem(BaseModel): + """Five members, each with a default the model answers for.""" + + nbnd: int = 1 + nosym: bool = False + ecutwfc: float = 60.0 + occupations: str = 'fixed' + degauss: float = 0.0 + + +class EngineRoute(BaseModel): + spin: Spin = Spin.NONE + system: EngineSystem = EngineSystem() + + +@task(input_model=EngineRoute, outputs=['report']) +def reports_its_system(spin, system): + return {'report': f'{spin.value}|{sorted(system)}'} + + +def test_the_leaf_body_is_handed_the_members_that_were_written(): + """One member written, and the body reads that member alone out of storage. + + The top-level field is the other half of the rule: nobody wrote ``spin``, + so the model answers for it and the body is handed its default. + """ + wg = WorkGraph('written_members') + node = wg.add_task(reports_its_system, name='leaf', system={'nbnd': 20}) + wg.run() + assert node.process.exit_status == 0 + assert node.outputs.report.value.value == "none|['nbnd']" + + +class EngineBlock(BaseModel): + """One member of a mapping, with defaults of its own.""" + + num_iter: int = 100 + dis_froz_max: float = 0.0 + num_wann: int = 4 + + +class EngineBlocks(BaseModel): + blocks: dict[str, EngineBlock] = {} + + +@task(input_model=EngineBlocks, outputs=['report']) +def reports_its_blocks(blocks): + return {'report': '|'.join(f'{key}:{sorted(item)}' for key, item in sorted(blocks.items()))} + + +def test_each_item_of_a_mapping_reaches_the_body_as_what_was_written(): + """A mapping is one namespace per key, so a key carries only its written members.""" + wg = WorkGraph('written_items') + node = wg.add_task(reports_its_blocks, name='leaf', blocks={'occ_1': {'num_iter': 42}}) + wg.run() + assert node.process.exit_status == 0 + assert node.outputs.report.value.value == "occ_1:['num_iter']" + + +# -------------------------------------------------------------------------- +# 13. What a calcfunction's parameters can be, and what a broken rule leaves +# -------------------------------------------------------------------------- + + +class NamespaceInputs(BaseModel): + """A parameter a process function has no port for.""" + + system: EngineSystem = EngineSystem() + + +class WindowInSteps(BaseModel): + """Two nodes and a rule that cannot answer until both are in hand.""" + + model_config = ConfigDict(arbitrary_types_allowed=True) + + lower: orm.Int + upper: orm.Int + + @model_validator(mode='after') + def _ordered(self): + if self.upper.value <= self.lower.value: + raise ValueError('upper must be above lower') + return self + + +@task.calcfunction(input_model=WindowInSteps) +def width(lower, upper): + return orm.Int(upper.value - lower.value) + + +def test_a_namespace_field_is_refused_on_a_calcfunction(): + """A calcfunction parameter is one port carrying one node, never a namespace.""" + with pytest.raises(ModelContractError, match='one port carrying one node'): + + @task.calcfunction(input_model=NamespaceInputs) + def reads_it(system): + return orm.Str(repr(sorted(system))) + + +def test_the_same_namespace_is_what_a_pyfunction_body_reads(): + """The control, and where the rule sends such a field.""" + wg = WorkGraph('namespace_pyfunction') + node = wg.add_task(reports_its_system, name='leaf', system={'nbnd': 20}) + wg.run() + assert node.process.exit_status == 0 + + +def test_a_calcfunction_runs_what_its_rule_admits(): + """The scalar-node shape a calcfunction does carry, rule and all.""" + wg = WorkGraph('width_ok') + node = wg.add_task(width, name='t', lower=orm.Int(1).store(), upper=orm.Int(9).store()) + wg.run() + assert node.process.exit_status == 0 + assert node.outputs.result.value.value == 8 + + +def test_a_rule_broken_inside_a_calcfunction_excepts_it(): + """What a calcfunction can carry when a rule fails: an exception, not an exit status. + + A process function has no controlled-failure channel for an exception in + its body, so the model's report reaches the excepted ``CalcFunctionNode`` + rather than an exit message. + """ + wg = WorkGraph('width_broken') + wg.add_task(width, name='t', lower=orm.Int(9).store(), upper=orm.Int(3).store()) + wg.run() + excepted = ( + orm.QueryBuilder() + .append( + orm.CalcFunctionNode, + tag='n', + filters={'label': 'width', 'attributes.process_state': 'excepted'}, + ) + .order_by({'n': {'id': 'desc'}}) + .first()[0] + ) + assert excepted.exit_status is None + assert 'upper must be above lower' in (excepted.exception or '') diff --git a/tests/test_serializer.py b/tests/test_serializer.py index 33f0b184..ea3d50c6 100644 --- a/tests/test_serializer.py +++ b/tests/test_serializer.py @@ -1,6 +1,11 @@ +from collections import OrderedDict, defaultdict, namedtuple +from enum import Enum, IntEnum + from aiida_workgraph import WorkGraph, task import pytest +from aiida_workgraph.serialization import AiidaSerializationAdapter, _flatten_enums + @task.graph() def sub_workflow(func): @@ -14,3 +19,230 @@ def test_func_as_input(capsys): wg.add_task(sub_workflow, func=add, name='sub_workflow') with pytest.raises(Exception, match='Cannot serialize the provided object'): wg.save() + + +class Color(Enum): + """Plain enum whose value is a bare string.""" + + RED = 'red' + BLUE = 'blue' + + +class Priority(IntEnum): + """Enum whose value is a bare int.""" + + LOW = 1 + HIGH = 9 + + +class Flavour(str, Enum): + """Enum member that is itself a ``str`` subclass.""" + + SWEET = 'sweet' + SOUR = 'sour' + + +def test_flatten_bare_enum_member(): + """A lone member collapses to its ``.value``, not a mangled ``Color.RED``.""" + assert _flatten_enums(Color.RED) == 'red' + assert not isinstance(_flatten_enums(Color.RED), Enum) + + +def test_flatten_intenum_and_str_enum_yield_bare_value(): + """The bare value's own type is preserved (int stays int, str stays str).""" + assert _flatten_enums(Priority.HIGH) == 9 + assert type(_flatten_enums(Priority.HIGH)) is int + assert _flatten_enums(Flavour.SWEET) == 'sweet' + assert type(_flatten_enums(Flavour.SWEET)) is str + + +def test_flatten_dict_keys_and_values(): + """Enums appearing as dict keys *and* values both flatten.""" + assert _flatten_enums({Color.RED: Color.BLUE, 'n': Priority.LOW}) == {'red': 'blue', 'n': 1} + + +def test_flatten_list_and_tuple_preserve_container_type(): + """Sequences recurse element-wise and keep list-vs-tuple identity.""" + assert _flatten_enums([Color.RED, 1, 'x']) == ['red', 1, 'x'] + out = _flatten_enums((Color.RED, Priority.HIGH)) + assert out == ('red', 9) + assert isinstance(out, tuple) + + +def test_flatten_mixed_nesting(): + """Enums buried in a dict/list/tuple mixture are all reached.""" + payload = { + 'a': [Color.RED, {'b': Priority.LOW}], + Color.BLUE: ('x', Flavour.SOUR), + } + assert _flatten_enums(payload) == { + 'a': ['red', {'b': 1}], + 'blue': ('x', 'sour'), + } + + +def test_flatten_passthrough_of_enum_free_payload(): + """An enum-free payload comes back with identical values.""" + payload = {'n': 3, 'items': [1, 'two', (3.0, None)], 'flag': True} + assert _flatten_enums(payload) == payload + + +def test_flatten_passthrough_preserves_namedtuple_type(): + """An enum-free namedtuple comes back as the exact same object, not + downgraded to a plain tuple.""" + Point = namedtuple('Point', ['x', 'y']) + p = Point(1, 2) + out = _flatten_enums(p) + assert out is p + assert isinstance(out, Point) + + +def test_flatten_namedtuple_with_enum_still_flattens(): + """A namedtuple carrying an Enum member still flattens its contents.""" + Point = namedtuple('Point', ['x', 'y']) + p = Point(Color.RED, 2) + out = _flatten_enums(p) + assert out == ('red', 2) + + +def test_flatten_passthrough_preserves_dict_subclass_type(): + """An enum-free OrderedDict/defaultdict comes back as the exact same + object, not downgraded to a plain dict.""" + od = OrderedDict([('a', 1), ('b', 2)]) + assert _flatten_enums(od) is od + + dd: defaultdict = defaultdict(int, {'a': 1}) + assert _flatten_enums(dd) is dd + + +def test_flatten_dict_subclass_with_enum_still_flattens(): + """An OrderedDict/defaultdict carrying an Enum still flattens its + values (into a plain dict -- rebuilding the original subclass is not + attempted for the changed case).""" + od = OrderedDict([('a', Color.RED), ('b', 2)]) + assert _flatten_enums(od) == {'a': 'red', 'b': 2} + + dd: defaultdict = defaultdict(int, {'a': Priority.LOW}) + assert _flatten_enums(dd) == {'a': 1} + + +def test_flatten_wrapt_proxied_enum_member(): + """A wrapt ``ObjectProxy`` (as TaggedValue wraps members) still flattens.""" + wrapt = pytest.importorskip('wrapt') + proxied = wrapt.ObjectProxy(Color.RED) + assert _flatten_enums(proxied) == 'red' + assert _flatten_enums({proxied: proxied}) == {'red': 'red'} + + +def test_flatten_dict_key_collision_raises(): + """Two distinct keys collapsing to the same flattened key raises rather than + silently dropping an entry. Here ``Color.RED`` (a plain Enum, so distinct in + identity and hash from the bare string) and ``'red'`` are separate keys that + both flatten to ``'red'``.""" + payload = {Color.RED: 1, 'red': 2} + assert len(payload) == 2 # distinct keys before flattening + with pytest.raises(ValueError, match='collision'): + _flatten_enums(payload) + + +def test_flatten_leaves_sets_untouched(): + """``set``/``frozenset`` are not descended into: they fail in + ``general_serializer`` regardless of contents, so flattening enums inside + them would not make them serializable. Pin that they pass through unchanged + (still holding Enum members).""" + s = {Color.RED, Color.BLUE} + out = _flatten_enums(s) + assert out is s + assert Color.RED in out + fs = frozenset({Priority.LOW}) + assert _flatten_enums(fs) is fs + + +def test_serialize_ports_accepts_enum_value(aiida_profile): + """End-to-end through the adapter: an enum-valued namespace entry serializes + without raising, landing as the flattened bare value (needs a profile so the + downstream ``general_serializer`` can build the node).""" + from aiida import orm + from node_graph.socket_spec import SocketMeta, SocketSpec + + spec = SocketSpec(identifier='node_graph.namespace', meta=SocketMeta(dynamic=True)) + out = AiidaSerializationAdapter().serialize_ports({'c': Color.RED, 'n': 3}, spec, store=False) + assert isinstance(out['c'], orm.Str) + assert out['c'].value == 'red' + + +@task(outputs=['type_name', 'is_member', 'equals_member', 'equals_value']) +def observe_enum(c: Color) -> dict: + """Report what a function task's body receives for an ``Enum``-typed + input, using only the annotation's contract (``==``), never a + reconstruction of the member.""" + return { + 'type_name': type(c).__name__, + 'is_member': isinstance(c, Color), + 'equals_member': c == Color.RED, + 'equals_value': c == 'red', + } + + +@task() +def echo_str(s: str) -> str: + return s + + +@task.graph() +def observe_enum_in_graph(c: Color) -> str: + """Use the member a ``@task.graph`` body receives, as the annotation + declares it — no reconstruction.""" + return echo_str(s=c.name) + + +def _node_graph_rebuilds_enums() -> bool: + """Return whether the installed ``node_graph`` reconstructs Enum sockets. + + ``coerce_inputs_from_spec`` rebuilds a socket's value only from the + ``structured_type`` descriptor its spec records, so a descriptor for an + ``Enum`` type is exactly the capability. + """ + from node_graph.utils.struct_utils import structured_type_info + + return structured_type_info(Color) is not None + + +def test_enum_arrival_follows_the_node_graph_capability(aiida_profile): + """A task body never reconstructs an ``Enum`` member by hand; the + annotation is the contract, and what arrives behind it is the installed + ``node_graph``'s call, not this package's. With a ``node_graph`` that + rebuilds ``Enum`` sockets (``_node_graph_rebuilds_enums``), a + ``Color``-annotated input arrives as the member: ``isinstance`` holds, + ``==`` against the member holds, ``type_name`` is ``'Color'``. Without + that capability it arrives as the flattened bare value: ``isinstance`` + fails, ``==`` against the value holds, ``type_name`` is ``'str'``.""" + wg = WorkGraph('enum_arrival') + t = wg.add_task(observe_enum, name='observe', c=Color.RED) + wg.run() + assert wg.state == 'FINISHED' + rebuilds = _node_graph_rebuilds_enums() + assert t.outputs['is_member'].value.value is rebuilds + assert t.outputs['type_name'].value.value == ('Color' if rebuilds else 'str') + if rebuilds: + assert t.outputs['equals_member'].value.value is True + else: + assert t.outputs['equals_value'].value.value is True + + +def test_enum_input_arrives_as_the_member_in_a_graph_body(aiida_profile): + """Only under a ``node_graph`` that rebuilds ``Enum`` sockets + (``_node_graph_rebuilds_enums``) can a ``@task.graph`` body use the + member directly, as its annotation declares (``c.name``, no + reconstruction).""" + if not _node_graph_rebuilds_enums(): + pytest.skip( + 'node_graph does not reconstruct Enum sockets ' + '(structured_type_info(Color) is None); ' + 'observe_enum_in_graph requires the member to use c.name.' + ) + wg = WorkGraph('enum_arrival_in_graph') + t = wg.add_task(observe_enum_in_graph, name='observe_graph', c=Color.RED) + wg.run() + assert wg.state == 'FINISHED' + assert t.outputs['result'].value.value == 'RED'