diff --git a/src/node_graph/decorator.py b/src/node_graph/decorator.py index 3bfc4b3..40ac5ac 100644 --- a/src/node_graph/decorator.py +++ b/src/node_graph/decorator.py @@ -1,6 +1,7 @@ from __future__ import annotations -from typing import Any, Optional, Callable, List, Dict +from typing import Any, Optional, Callable, List, Dict, Type import inspect +from pydantic import BaseModel from .error_handler import ErrorHandlerSpec from .task import Task from .task_spec import TaskHandle, BaseHandle @@ -38,6 +39,8 @@ def decorator_task( outputs: Optional[SocketSpec | List[str]] = None, error_handlers: Optional[Dict[str, ErrorHandlerSpec]] = None, catalog: str = "Others", + input_model: Optional[Type[BaseModel]] = None, + output_model: Optional[Type[BaseModel]] = None, ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: """Generate a decorator that register a function as a Graph task. After decoration, calling that function `func(x, y, ...)` @@ -49,11 +52,19 @@ def decorator_task( catalog (str): task catalog inputs (dict): task inputs outputs (dict): 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 wrap(func) -> TaskHandle: from node_graph.tasks.function_task import FunctionTask + from node_graph.input_model import apply_models + in_spec, out_spec, executor = apply_models( + func, inputs, outputs, input_model, output_model + ) callable_meta = inspect_callable_metadata(func) metadata = {"callable": callable_meta} version = callable_meta.get("package_version") @@ -69,11 +80,12 @@ def wrap(func) -> TaskHandle: obj=func, identifier=resolved_identifier, catalog=catalog, - input_spec=inputs, - output_spec=outputs, + input_spec=in_spec, + output_spec=out_spec, error_handlers=error_handlers, metadata=metadata, version=version, + executor=executor, ) return wrap @@ -84,6 +96,8 @@ def decorator_graph( inputs: Optional[SocketSpec | list] = None, outputs: Optional[SocketSpec | list] = None, catalog: str = "Others", + input_model: Optional[Type[BaseModel]] = None, + output_model: Optional[Type[BaseModel]] = None, ) -> Callable[[Callable[..., Any]], Callable[..., Any]]: """Generate a decorator that register a function as a graph task. @@ -92,11 +106,19 @@ def decorator_graph( catalog (str): task catalog inputs (dict): task inputs outputs (dict): 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 wrap(func) -> TaskHandle: from node_graph.tasks.function_task import FunctionTask + from node_graph.input_model import apply_models + in_spec, _, executor = apply_models( + func, inputs, None, input_model, output_model, is_graph=True + ) callable_meta = inspect_callable_metadata(func) metadata = {"callable": callable_meta} version = callable_meta.get("package_version") @@ -113,10 +135,11 @@ def wrap(func) -> TaskHandle: identifier=resolved_identifier, task_type="graph", catalog=catalog, - input_spec=inputs, + input_spec=in_spec, output_spec=outputs, metadata=metadata, version=version, + executor=executor, ) return wrap diff --git a/src/node_graph/errors.py b/src/node_graph/errors.py index 471c352..4f1a20c 100644 --- a/src/node_graph/errors.py +++ b/src/node_graph/errors.py @@ -2,3 +2,24 @@ class GraphDeferredIllegalOperationError(TypeError): """Raised when a future (Socket) is used in an illegal concrete operation.""" pass + + +class SocketValueError(ValueError): + """A value the socket layer refused, carrying where it was refused. + + ``loc`` names the socket path below a task's inputs and ``type`` names the + rule that refused, both spelled as pydantic spells them in + ``ValidationError.errors()``: a caller reading a refusal does not have to + know whether an input model or the socket layer produced it. + """ + + def __init__( + self, + message: str, + *, + loc: tuple = (), + error_type: str = "", + ) -> None: + super().__init__(message) + self.loc = tuple(loc) + self.type = error_type diff --git a/src/node_graph/executor.py b/src/node_graph/executor.py index 23fe3b7..912616f 100644 --- a/src/node_graph/executor.py +++ b/src/node_graph/executor.py @@ -17,6 +17,10 @@ class ExecutorMode(str, Enum): PICKLED_CALLABLE = "pickled_callable" # A callable serialized using cloudpickle. +#: Attribute holding the name a callable's module binds it under. +BOUND_NAME_ATTR = "__node_graph_bound_name__" + + def serialize_callable( func: Callable, register_pickle_by_value: bool = False, include_source: bool = True ) -> Dict[str, Any]: @@ -43,7 +47,10 @@ def serialize_callable( if not callable(func): raise TypeError("Provided object is not a callable function or class.") source_code = "" - callable_name = func.__name__ + # A wrapper its module binds under a name of its own says so, because + # ``functools.wraps`` gives it the wrapped function's ``__name__`` and + # importing that name back would return the function, not the wrapper. + callable_name = getattr(func, BOUND_NAME_ATTR, None) or func.__name__ if func.__module__ == "__main__" or "." in func.__qualname__.split(".", 1)[-1]: mode = ExecutorMode.PICKLED_CALLABLE pickled_data = cloudpickle.dumps(func) diff --git a/src/node_graph/input_model.py b/src/node_graph/input_model.py new file mode 100644 index 0000000..86e3e46 --- /dev/null +++ b/src/node_graph/input_model.py @@ -0,0 +1,1944 @@ +"""Make a Pydantic model the wire contract of a task. + +``@task(input_model=M)`` builds the task's input sockets from ``M`` and +enforces ``M`` at three moments. The body keeps a plain-Python signature and +receives plain Python: ``M`` decides types, defaults and cross-field rules; +the signature only names the parameters. + +Three checkpoints, three questions: + +- **wiring** (:func:`validate_wiring_inputs`): every value written must + satisfy the field it is written to. The model runs as a flat shadow, so a + field's *type* is checked and a socket or a task written into an input + passes untouched -- what flows through it is not known yet. The model's own + ``mode='after'`` ``@field_validator``s then run over the fields whose value + is resolved: a tagged value is judged by the value it carries, and a field + holding a socket or a task waits, as does one rule reaching through + ``info.data`` for a field the write did not name. A ``@model_validator`` + waits in every case -- the model's own and those of the models its fields + declare -- because it may read a field no one has written yet. +- **graph expansion** (:func:`validate_graph_inputs`): a ``@task.graph``'s + inputs are resolved by then, so the real model runs, cross-field rules + included, and the body receives what the model made of each field, wearing + the tag that field was written with -- the tag, because it is what the body + turns into a link. +- **run edge** (:func:`validated_callable`): a leaf task's inputs are + validated once the engine has assembled and deserialized them, and the body + receives the rich objects ``M`` declares. + +At the last two, validation may change a value's *representation* and never +its *content* (:func:`check_content_invariance`): a model that derives or +rewrites an input is refused, because the derived value would reach the body +while storage and provenance keep the value the caller wrote. + +Storage is the engine adapter's business, but the rendering is the model's: +:func:`model_dumper_for_socket` resolves a socket's path through the model +tree and returns the function that renders that leaf. Reading back is the +same bargain the other way round: every leaf socket records what a body +receives for it (:data:`BODY_RECEIVES`), so a field declaring a type pydantic +can rebuild from plain data arrives as plain Python, and one declaring an +engine's own data class -- or ``Any`` -- arrives as whatever the engine +stored. + +A mapping whose size is decided at runtime is written as a typed container +field, ``dict[str, T]``: it becomes a dynamic namespace whose members are +sockets named by the mapping's keys, each shaped by ``T``. It is also the +only shape that carries engine data classes: a list of them is refused, +because storage keeps a list as one node of plain data. ``extra='allow'`` is +refused at any depth -- a model that admits fields it does not declare is not +a contract. +""" + +from __future__ import annotations + +import functools +import inspect +import sys +from copy import copy +from dataclasses import MISSING, replace +from typing import ( + Annotated, + Any, + Callable, + Dict, + FrozenSet, + List, + Literal, + NamedTuple, + Optional, + Tuple, + Type, + cast, + get_args, + get_origin, + get_type_hints, +) + +from pydantic import ( + BaseModel, + ConfigDict, + PlainSerializer, + TypeAdapter, + ValidationError, + WrapSerializer, + WrapValidator, + create_model, + field_validator, +) +from pydantic.errors import PydanticSchemaGenerationError +from pydantic_core import PydanticSerializationError +from typing_extensions import is_typeddict + +__all__ = [ + "ModelContractError", + "ModelDerivedValueError", + "TaskInputValidationError", + "TaskOutputValidationError", + "BODY_RECEIVES", + "ContentSnapshot", + "apply_models", + "body_inputs", + "check_content_invariance", + "content_snapshot", + "check_signature_against_model", + "dump_model_field", + "input_model_of_callable", + "model_dumper_for_socket", + "rebind_executor_callable", + "spec_from_model", + "validate_graph_inputs", + "validate_task_inputs", + "validate_wiring_inputs", + "validated_callable", +] + +#: Attributes under which a callable carries the models its task enforces. +INPUT_MODEL_ATTR = "__node_graph_input_model__" +OUTPUT_MODEL_ATTR = "__node_graph_output_model__" + +#: Attribute caching a task's input model, so the executor is resolved once. +_TASK_MODEL_CACHE = "_node_graph_cached_input_model" + +#: Sentinel telling "no model" apart from "not looked up yet". +_UNRESOLVED = object() + +#: Sentinel standing for a field no partial write named. +_MISSING = object() + +#: How deep a generic annotation is rebuilt before it is taken as a leaf. +_MAX_ANNOTATION_DEPTH = 6 + +#: ``Optional[T]``'s second argument, which stands for itself in a rebuild. +_NONE_TYPE = type(None) + +#: Socket extra naming what a body receives for a leaf: ``"python"`` or ``"node"``. +BODY_RECEIVES = "body_receives" + +#: Container origins whose members storage flattens into one node. +_SEQUENCE_ORIGINS = (list, tuple, set, frozenset) + +#: Annotation metadata that renders a value, which the content check must not see. +_SERIALIZER_METADATA = (PlainSerializer, WrapSerializer) + +#: ``populate_by_name``'s replacement, on pydantic 2.11 and later. +_HAS_VALIDATE_BY_NAME = "validate_by_name" in getattr(ConfigDict, "__annotations__", {}) + + +class ModelContractError(TypeError): + """Raised when a model and the function it describes disagree.""" + + +class TaskInputValidationError(ValueError): + """Raised when a task's assembled inputs fail its input model. + + ``task`` names what the value was written to and ``model`` is the model + that refused it. ``errors`` is pydantic's own report, one entry per + refused value carrying the ``loc``, ``type`` and ``msg`` pydantic spells + them with, so a caller reads a refusal rather than parsing its message -- + the same two facts :class:`~node_graph.errors.SocketValueError` carries + for a value the socket layer refuses. It is empty when the refusal was + raised without a pydantic report behind it. + """ + + def __init__( + self, + message: str, + *, + task: str = "", + model: Optional[Type[BaseModel]] = None, + errors: Optional[List[Dict[str, Any]]] = None, + ) -> None: + super().__init__(message) + self.task = task + self.model = model + self.errors: List[Dict[str, Any]] = list(errors or ()) + + +class TaskOutputValidationError(ValueError): + """Raised when a task's return value fails its output model.""" + + +class ModelDerivedValueError(ValueError): + """Raised when validating an input changes its content.""" + + +def _is_model(obj: Any) -> bool: + return isinstance(obj, type) and issubclass(obj, BaseModel) + + +# -------------------------------------------------------------------------- +# The socket spec a model describes +# -------------------------------------------------------------------------- + + +def _socket_api(api: Any = None) -> Any: + """Return the ``SocketSpecAPI`` building the spec, node-graph's by default.""" + if api is not None: + return api + from node_graph.socket_spec import SocketSpecAPI + + return SocketSpecAPI + + +def _strip_structured_types(spec: Any) -> Any: + """Drop the ``structured_type`` descriptor from every socket in ``spec``. + + The descriptor names a class by dotted import path and is persisted with + the task, so a class that moves or is renamed breaks stored data. Under a + model contract nothing needs it: the model reconstructs rich values at the + run edge, and the model itself is reached through the task's executor + rather than through a path recorded in the data. + """ + meta = spec.meta + if "structured_type" in (meta.extras or {}): + extras = dict(meta.extras) + extras.pop("structured_type") + spec = replace(spec, meta=replace(meta, extras=extras)) + if spec.fields: + spec = replace( + spec, + fields={ + name: _strip_structured_types(child) + for name, child in spec.fields.items() + }, + ) + if spec.item is not None: + spec = replace(spec, item=_strip_structured_types(spec.item)) + return spec + + +def _strip_optional_type(annotation: Any) -> Any: + from node_graph.socket_spec import _strip_optional + + return _strip_optional(annotation) + + +def _container_item_type(annotation: Any) -> Any: + """Return the value type of a ``dict[str, T]`` field, else ``None``. + + A bare ``dict`` names no value type and stays a leaf. ``list[T]`` is not a + container in this sense: a namespace addresses its members by name, and a + list has none to give. + """ + base = _strip_optional_type(annotation) + if get_origin(base) is not dict: + return None + args = get_args(base) + if not args: + return None + key_type, item_type = args + if key_type not in (str, Any): + raise ModelContractError( + f"A dynamic field is keyed by socket name, so its keys must be str; got {key_type!r}.\n" + "How to fix: declare the field as dict[str, ...]." + ) + return item_type + + +def _model_of_type(annotation: Any) -> Optional[Type[BaseModel]]: + """Return the model an annotation resolves to, if it resolves to one.""" + base = _strip_optional_type(annotation) + return base if _is_model(base) else None + + +def _union_arms(annotation: Any) -> Optional[List[Any]]: + """Return a union's arms other than ``None``, or ``None`` if it is no union.""" + from node_graph.socket_spec import _is_union_origin + + if not _is_union_origin(get_origin(annotation)): + return None + return [arg for arg in get_args(annotation) if arg is not _NONE_TYPE] + + +def _annotation_name(annotation: Any) -> str: + """Render an annotation the way it is written.""" + if isinstance(annotation, type) and get_origin(annotation) is None: + return annotation.__name__ + return repr(annotation) + + +def _plural(items: List[str], singular: str, plural: str) -> str: + """Return the word ``items``'s number takes.""" + return singular if len(items) == 1 else plural + + +def _declares_any(annotation: Any) -> bool: + """Return True when ``annotation`` is ``Any``, under an ``Optional`` or not.""" + return _strip_optional_type(annotation) is Any + + +def _disagreeing_kinds( + label: str, + subject: str, + parts: str, + kinds: List[Tuple[str, str]], + *, + declares_any: bool, +) -> ModelContractError: + """Return the error raised when the pieces of a type arrive in both forms. + + ``kinds`` pairs each piece's rendered name with ``"python"`` or ``"node"``. + ``declares_any`` says whether one of the pieces is ``Any``, which is worth + naming: it is on the stored side because it declares nothing to rebuild. + """ + rebuilt = [name for name, kind in kinds if kind == "python"] + stored = [name for name, kind in kinds if kind == "node"] + note = " (Any declares nothing to rebuild)" if declares_any else "" + return ModelContractError( + f"{label} declares {subject}, whose {parts} disagree on how the value arrives: " + f"{', '.join(rebuilt)} {_plural(rebuilt, 'is', 'are')} rebuilt from plain data, " + f"{', '.join(stored)} {_plural(stored, 'arrives', 'arrive')} as the engine stored " + f"{_plural(stored, 'it', 'them')}{note}.\n" + f"How to fix: declare {label} so its {parts} agree. A socket delivers one form, so a " + f"value spanning both leaves nothing to decide it." + ) + + +def _declared_members(annotation: Any) -> Optional[List[Tuple[str, Any]]]: + """Return the ``(name, type)`` members ``annotation`` declares, else ``None``. + + A model and a ``TypedDict`` name their members' types, and those decide + how the value arrives: pydantic builds the mapping around an ``orm.Int`` + readily enough, and the ``orm.Int`` inside it still only an instance + satisfies. A container is followed only as far as a member that declares + members of its own, so ``list[int]`` is left to the plain reading below. + """ + base = _strip_optional_type(annotation) + if _is_model(base): + return [(name, field.annotation) for name, field in base.model_fields.items()] + if is_typeddict(base): + try: + hints = get_type_hints(base, include_extras=True) + except Exception: + return None + return list(hints.items()) + origin = get_origin(base) + if origin in _SEQUENCE_ORIGINS or origin is dict: + args = list(get_args(base)) + if origin is dict: + args = args[1:] + members = [ + (f"[{index}]", arg) + for index, arg in enumerate(args) + if arg is not Ellipsis + and arg is not _NONE_TYPE + and _declared_members(arg) is not None + ] + return members or None + return None + + +def _refuse_container_of_nodes(base: Any, label: str, depth: int) -> None: + """Raise when a sequence's members are values only an instance satisfies. + + Storage keeps a sequence as one node holding plain data, so its members + come back as plain data whatever they were written as. A mapping is the + shape that survives: ``dict[str, T]`` becomes one socket, and one node, + per key. + """ + if get_origin(base) not in _SEQUENCE_ORIGINS: + return + for arg in get_args(base): + if arg is Ellipsis or arg is _NONE_TYPE or _strip_optional_type(arg) is Any: + continue + if _body_receives(arg, f"{label}[]", depth + 1) != "node": + continue + name = _annotation_name(_strip_optional_type(arg)) + raise ModelContractError( + f"{label} declares {_annotation_name(base)}, and a container of {name} cannot " + f"round-trip through storage, which keeps it as one node holding plain data.\n" + f"How to fix: declare {label} as dict[str, {name}], which gives each member a " + f"socket and a node of its own, or as a container of plain data." + ) + + +def _arm_receives(annotation: Any, label: str, depth: int = 0) -> str: + """Return ``"python"`` or ``"node"`` for one type, unions not considered.""" + base = _strip_optional_type(annotation) + if base is Any or base is inspect.Parameter.empty: + return "node" + if depth <= _MAX_ANNOTATION_DEPTH: + _refuse_container_of_nodes(base, label, depth) + members = _declared_members(base) + if members is not None: + kinds = [ + (name, _body_receives(member, f"{label}.{name}", depth + 1)) + for name, member in members + ] + answers = {kind for _, kind in kinds} + if len(answers) == 1: + return answers.pop() + raise _disagreeing_kinds( + label, + _annotation_name(base), + "members", + kinds, + declares_any=any(_declares_any(member) for _, member in members), + ) + try: + TypeAdapter(annotation, config=ConfigDict(arbitrary_types_allowed=False)) + except PydanticSchemaGenerationError: + return "node" + except Exception: + # A dataclass or TypedDict refuses the config argument because it + # carries its own; each is rebuilt from plain data all the same. + return "python" + return "python" + + +def _body_receives(annotation: Any, label: str, depth: int = 0) -> str: + """Return ``"python"`` or ``"node"`` for what a body is handed for ``annotation``. + + ``"python"`` when pydantic can build the declared type out of plain data, + so the engine's storage wrapper has to come off before the model sees the + value. ``"node"`` when the field can only be satisfied by an instance of a + class pydantic treats as arbitrary -- an engine's own data class, say -- + and when the field declares ``Any``, which declares nothing to rebuild. + + A type that declares members of its own -- a nested model, a ``TypedDict`` + -- is read through them, because they are what has to be satisfied. + + A union is read one arm at a time and every arm must give the same answer. + ``X | None`` is not a union in this sense: ``None`` is neither, so such a + field takes ``X``'s answer. A union spanning both -- ``int | orm.Int`` -- + is refused, and so are members that disagree: a socket delivers one form, + and nothing here can choose. + """ + arms = _union_arms(annotation) + if arms is None or len(arms) == 1: + return _arm_receives(annotation if arms is None else arms[0], label, depth) + kinds = [(arm, _arm_receives(arm, label, depth)) for arm in arms] + answers = {kind for _, kind in kinds} + if len(answers) == 1: + return answers.pop() + raise _disagreeing_kinds( + label, + repr(annotation), + "arms", + [(_annotation_name(arm), kind) for arm, kind in kinds], + declares_any=any(_declares_any(arm) for arm, _ in kinds), + ) + + +def _mark_body_arrival(spec: Any, annotation: Any, label: str) -> Any: + """Record on a leaf socket what a body receives for it. + + A namespace holds no value of its own, so only leaves are marked. The mark + travels with the stored spec, which is what lets the read edge answer for a + socket alone, without resolving the model behind it. + """ + if spec.is_namespace(): + return spec + extras = dict(spec.meta.extras or {}) + extras[BODY_RECEIVES] = _body_receives(annotation, label) + return replace(spec, meta=replace(spec.meta, extras=extras)) + + +def _fields_from_model( + model: Type[BaseModel], spec: Any, api: Any, *, defaults: bool = True +) -> Any: + """Overlay on ``spec`` what the model knows and ``from_model`` does not. + + Four things: a field's requiredness, which ``from_model`` leaves at + ``True`` whatever the field's default; a default rendered JSON-safe, + because the spec is stored with the task; a ``dict[str, T]`` field turned + into a typed dynamic namespace, so each key of the mapping is a socket of + its own and can be linked by name; and, on every leaf, what a body + receives for it (:func:`_body_receives`). + + ``defaults`` is false below a nested model and below a mapping's item + model: a member the caller left out carries no default onto its socket, + so nothing is collected for it and the body is handed the members that + were written. The member stays optional, so a write may still name any + subset of them. + """ + fields = dict(spec.fields or {}) + for name, field in model.model_fields.items(): + child = fields.get(name) + if child is None: + continue + item_type = _container_item_type(field.annotation) + if item_type is not None: + child = api.dynamic(item_type) + item_model = _model_of_type(item_type) + if item_model is not None and child.item is not None: + child = replace( + child, + item=_fields_from_model( + item_model, child.item, api, defaults=False + ), + ) + elif child.item is not None: + child = replace( + child, + item=_mark_body_arrival( + child.item, item_type, f"{model.__name__}.{name}[]" + ), + ) + else: + nested = _model_of_type(field.annotation) + if nested is not None and child.is_namespace(): + child = _fields_from_model(nested, child, api, defaults=False) + else: + child = _mark_body_arrival( + child, field.annotation, f"{model.__name__}.{name}" + ) + child = replace(child, meta=replace(child.meta, required=field.is_required())) + if child.is_namespace(): + pass + elif not defaults: + child = replace(child, default=MISSING) + elif not field.is_required(): + child = replace( + child, + default=dump_model_field( + model, name, field.get_default(call_default_factory=True) + ), + ) + fields[name] = child + return replace(spec, fields=fields) + + +def _models_reached_by(annotation: Any, depth: int = 0) -> List[Type[BaseModel]]: + """Return every model ``annotation`` reaches, through containers and unions.""" + if depth > _MAX_ANNOTATION_DEPTH: + return [] + base = _strip_optional_type(annotation) + if _is_model(base): + return [base] + found: List[Type[BaseModel]] = [] + for arg in get_args(base): + found.extend(_models_reached_by(arg, depth + 1)) + return found + + +def _refuse_open_models(model: Type[BaseModel], seen: set, depth: int = 0) -> None: + """Raise for any model reachable from ``model`` that admits undeclared fields. + + A nested model is checked too: a field it admits but does not declare has + no socket, so it reaches the body without ever reaching storage. + """ + if model in seen or depth > _MAX_ANNOTATION_DEPTH: + return + seen.add(model) + if model.model_config.get("extra") == "allow": + raise ModelContractError( + f"{model.__name__} sets extra='allow', which declares no contract for the fields it admits.\n" + "How to fix: declare the dynamic part as a field of type dict[str, T]." + ) + for field in model.model_fields.values(): + for nested in _models_reached_by(field.annotation): + _refuse_open_models(nested, seen, depth + 1) + + +def _refuse_models_that_name_themselves( + model: Type[BaseModel], path: Tuple[Type[BaseModel], ...] = () +) -> None: + """Raise when a model reachable from ``model`` is reachable from itself. + + A namespace holds one socket per field, so a model whose field names a + model already on the way to it describes a namespace with no bottom. + """ + if model in path: + cycle = " -> ".join(item.__name__ for item in (*path, model)) + raise ModelContractError( + f"{cycle} names a model already on the way to it, and a socket " + "namespace has to end somewhere.\n" + "How to fix: give the recursive part its own task, or declare the " + "field as plain data the body reads for itself." + ) + for field in model.model_fields.values(): + for nested in _models_reached_by(field.annotation): + _refuse_models_that_name_themselves(nested, (*path, model)) + + +def spec_from_model(model: Type[BaseModel], api: Any = None) -> Any: + """Return the socket namespace ``model`` describes. + + An open-topped model is refused, at any depth: ``extra='allow'`` accepts + fields nothing declared, which is the one shape a contract cannot + describe. So is a model that names itself, which describes a namespace + with no bottom. A mapping whose size is only known at runtime is written + as a typed container field, ``dict[str, T]``. + """ + api = _socket_api(api) + _refuse_open_models(model, set()) + _refuse_models_that_name_themselves(model) + return _strip_structured_types( + _fields_from_model(model, api.from_model(model), api) + ) + + +# -------------------------------------------------------------------------- +# The model and the function it describes +# -------------------------------------------------------------------------- + + +def _annotations_agree(declared: Any, annotation: Any) -> bool: + """Return True when a signature annotation may stand for a field's type.""" + if annotation is inspect.Parameter.empty or annotation is Any: + return True + return bool(annotation == declared) + + +def check_signature_against_model( + func: Callable[..., Any], + model: Type[BaseModel], + *, + role: str, +) -> None: + """Raise unless every field of ``model`` names a parameter of ``func``. + + ``role`` is ``'input'`` or ``'output'``; an output model is checked only + for the shape a model may take, since it describes the return value rather + than parameters. + """ + name = getattr(func, "__name__", repr(func)) + if not _is_model(model): + raise ModelContractError( + f"{role}_model for '{name}' must be a pydantic BaseModel subclass; got {model!r}." + ) + if role == "output": + return + + signature = inspect.signature(func) + for parameter in signature.parameters.values(): + if parameter.kind in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + ): + raise ModelContractError( + f"'{name}' takes {parameter.name} as *args/**kwargs, which no model field can name.\n" + f"How to fix: list the parameters explicitly, one per field of {model.__name__}." + ) + if parameter.kind is inspect.Parameter.POSITIONAL_ONLY: + raise ModelContractError( + f"Parameter '{parameter.name}' of '{name}' is positional-only; a model field is passed by name.\n" + "How to fix: move it after the / marker, or drop the marker." + ) + + parameters = dict(signature.parameters) + declared = model.model_fields + + unnamed = [field for field in declared if field not in parameters] + if unnamed: + raise ModelContractError( + f"{model.__name__} declares {', '.join(repr(f) for f in unnamed)}, which '{name}' does not take.\n" + f"How to fix: add the parameter(s) to {name}, or remove the field(s) from {model.__name__}." + ) + + uncovered = [parameter for parameter in parameters if parameter not in declared] + if uncovered: + raise ModelContractError( + f"'{name}' takes {', '.join(repr(p) for p in uncovered)}, which {model.__name__} does not declare.\n" + f"How to fix: add the field(s) to {model.__name__}, or remove the parameter(s) from {name}." + ) + + defaulted = [ + parameter.name + for parameter in parameters.values() + if parameter.default is not inspect.Parameter.empty + ] + if defaulted: + raise ModelContractError( + f"'{name}' gives {', '.join(repr(p) for p in defaulted)} a default in its signature.\n" + f"How to fix: defaults live in the model - move it to {model.__name__} and leave the parameter bare." + ) + + try: + hints = inspect.get_annotations(func, eval_str=True) + except Exception: + hints = getattr(func, "__annotations__", {}) + for field_name, field in declared.items(): + annotation = hints.get(field_name, inspect.Parameter.empty) + if not _annotations_agree(field.annotation, annotation): + raise ModelContractError( + f"Parameter '{field_name}' of '{name}' is annotated {annotation!r}, " + f"but {model.__name__} declares {field.annotation!r}.\n" + "How to fix: drop the annotation, or make it match the field." + ) + + +# -------------------------------------------------------------------------- +# Checkpoint A -- the value written at the call +# -------------------------------------------------------------------------- + + +def is_socket_reference(value: Any) -> bool: + """Return True when ``value`` stands for something a link will deliver. + + A task counts: writing one into an input links its top-level output, so + what arrives is decided later, exactly as for a socket. A tag does not: + it carries the value it stands for, and that value is here now. + """ + from node_graph.socket import BaseSocket, TaggedValue + from node_graph.task import Task + + if isinstance(value, TaggedValue): + value = value.__wrapped__ + return isinstance(value, (BaseSocket, Task)) + + +def _awaits_a_link(value: Any, depth: int = 0) -> bool: + """Return True when ``value`` holds anything a link has yet to deliver. + + A tag is looked through: it carries the value it stands for, and a value + is what a rule can run on. A socket or a task is not -- what arrives + through it is decided later -- and one nested anywhere inside a mapping, a + sequence or a model instance leaves the whole value waiting. So does a + value nested deeper than the walk goes: what has not been looked at is not + known to have arrived. + """ + from node_graph.socket import BaseSocket, TaggedValue + from node_graph.task import Task + + if isinstance(value, TaggedValue): + value = value.__wrapped__ + if isinstance(value, (BaseSocket, Task)): + return True + if depth > _MAX_ANNOTATION_DEPTH: + return True + if isinstance(value, BaseModel): + return any( + _awaits_a_link(getattr(value, name, None), depth + 1) + for name in type(value).model_fields + ) + if isinstance(value, dict): + return any(_awaits_a_link(item, depth + 1) for item in value.values()) + if isinstance(value, _SEQUENCE_ORIGINS): + return any(_awaits_a_link(item, depth + 1) for item in value) + return False + + +def _accept_reference(value: Any, handler: Any, info: Any) -> Any: + """Let a socket reference through untouched; validate anything else. + + The test is :func:`is_socket_reference`, the one the rule pass uses too: + a field's type and its rules must judge the same write the same way, or a + value one of them refuses is a value the other never sees. + """ + from node_graph.socket import TaggedValue + + if is_socket_reference(value): + return value + if isinstance(value, TaggedValue): + value = value.__wrapped__ + return handler(value) + + +def _accepting_instances_of(model: Type[BaseModel]) -> Callable[..., Any]: + """Return a wrap validator that also lets an instance of ``model`` through. + + A nested field is checked against a shadow of the model it declares, and a + shadow is a different class: pydantic would refuse the very class the field + names. An instance already satisfies the declared type, so it passes as it + is, wherever it is written -- alone, or inside a ``list[M]`` or + ``dict[str, M]``. + """ + + def accept(value: Any, handler: Any, info: Any) -> Any: + from node_graph.socket import TaggedValue + + if is_socket_reference(value): + return value + if isinstance(value, TaggedValue): + value = value.__wrapped__ + if isinstance(value, model): + return value + return handler(value) + + return accept + + +def _rebuild_generic(annotation: Any, rebuild: Callable[[Any], Any]) -> Any: + """Return ``annotation`` with ``rebuild`` applied to each of its arguments. + + ``Literal``'s arguments are values rather than types and ``Annotated``'s + tail is metadata, so both keep their arguments; anything else that cannot + be rebuilt is returned unchanged. + """ + origin = get_origin(annotation) + args = get_args(annotation) + if origin is None or not args: + return None + if origin is Literal: + return annotation + if origin is Annotated: + return Annotated[tuple([rebuild(args[0]), *args[1:]])] + try: + return origin[ + tuple(rebuild(arg) if arg is not _NONE_TYPE else arg for arg in args) + ] + except TypeError: + return annotation + + +def _accepting_field_names(config: Dict[str, Any]) -> Dict[str, Any]: + """Return ``config`` with every field reachable by the name it is declared under.""" + config = dict(config) + config["populate_by_name"] = True + if _HAS_VALIDATE_BY_NAME: + config["validate_by_name"] = True + return config + + +def _shadow_config(model: Type[BaseModel]) -> ConfigDict: + """Return the config a model rebuilt from ``model``'s fields must carry. + + ``model_config`` decides what a field accepts and how it coerces -- + whitespace stripped, a class pydantic would otherwise refuse to build a + schema for. A rebuilt model that dropped it would judge the same value + differently from the model it stands for. + + One key is added rather than copied: a socket is named by its field, so + every rebuilt model accepts that name whatever alias the field carries. + """ + return cast(ConfigDict, _accepting_field_names(dict(model.model_config))) + + +@functools.lru_cache(maxsize=None) +def _by_field_name(model: Type[BaseModel]) -> Type[BaseModel]: + """Return ``model`` accepting each field under the name that names its socket. + + A field carrying ``Field(alias=...)`` is written by its alias and its + socket is named by the field, so the model as written would refuse the + very keys the graph delivers. The result is the model with one config key + added -- its rules, its types, its name in an error message. + """ + config = _accepting_field_names(dict(model.model_config)) + if config == dict(model.model_config): + return model + twin = type( + model.__name__, (model,), {"model_config": ConfigDict(**config)} # type: ignore[misc] + ) + twin.__qualname__ = model.__qualname__ + return cast(Type[BaseModel], twin) + + +def _reference_tolerant( + annotation: Any, depth: int = 0, *, partial: bool = False +) -> Any: + """Return ``annotation`` with every level willing to hold a reference. + + Wrapping only the outermost level is not enough: a reference written into + one member of a ``dict[str, int]`` reaches the ``int`` validator, which + refuses it. Every level is wrapped, so a reference is accepted wherever it + is written and every literal beside it is still checked. + + ``partial`` says the write names some of what will be written, and it + reaches every model below as well: a member of a nested namespace is as + free to arrive by a later link as a field of the model itself. + """ + if depth > _MAX_ANNOTATION_DEPTH: + return Annotated[annotation, WrapValidator(_accept_reference)] + rebuilt = _rebuild_generic( + annotation, lambda arg: _reference_tolerant(arg, depth + 1, partial=partial) + ) + if rebuilt is not None: + return Annotated[rebuilt, WrapValidator(_accept_reference)] + if _is_model(annotation): + shadow = ( + _partial_wiring_shadow(annotation, depth + 1) + if partial + else _wiring_shadow(annotation, depth + 1) + ) + return Annotated[shadow, WrapValidator(_accepting_instances_of(annotation))] + return Annotated[annotation, WrapValidator(_accept_reference)] + + +@functools.lru_cache(maxsize=None) +def _wiring_shadow(model: Type[BaseModel], depth: int = 0) -> Type[BaseModel]: + """Return the model checking types at the call, references included. + + The shadow is built with no base class, so the user's ``@field_validator`` + and ``@model_validator`` are left out: a rule written for whole, resolved + inputs would otherwise be judged at wiring time against a placeholder, and + a proxy that forwards comparisons makes that failure silent rather than + loud. It keeps ``model_config``, because the config decides what a field + accepts -- an alias reachable by field name, a class pydantic would + otherwise refuse to build a schema for -- and a shadow that judged those + differently would refuse calls the model accepts. + + ``depth`` counts the models the rebuild has walked through, so a model + that names itself terminates: past ``_MAX_ANNOTATION_DEPTH`` the field + keeps the class it declares, which is the same bound the resolved-value + walk stops at. + """ + fields = { + name: (_reference_tolerant(field.annotation, depth), field) + for name, field in model.model_fields.items() + } + return create_model( # type: ignore[call-overload] + f"{model.__name__}__Wiring", + __config__=_shadow_config(model), + **fields, + ) + + +def _optional_field(field: Any) -> Any: + """Return ``field`` free to be missing, keeping its type and constraints.""" + field = copy(field) + field.default = _MISSING + field.default_factory = None + field.validate_default = False + return field + + +@functools.lru_cache(maxsize=None) +def _partial_wiring_shadow(model: Type[BaseModel], depth: int = 0) -> Type[BaseModel]: + """Return the wiring shadow with every field, at every depth, free to be missing. + + Inputs may be written a few at a time and a link may supply the rest, so + a write that names some of the fields says nothing about the others -- + the members of a namespace it reaches into included. What it does say is + checked exactly as at a call: each value against the field it is written + to. + + ``depth`` bounds the rebuild where :func:`_wiring_shadow`'s does, so a + model that names itself terminates here too. + """ + fields = { + name: ( + _reference_tolerant(field.annotation, depth, partial=True), + _optional_field(field), + ) + for name, field in model.model_fields.items() + } + return create_model( # type: ignore[call-overload] + f"{model.__name__}__Partial", + __config__=_shadow_config(model), + **fields, + ) + + +class _UnwrittenSibling(Exception): + """A rule read a field the write it is judging did not name.""" + + +class _WrittenFields(dict): + """The fields a write named, refusing a lookup of any other. + + ``info.data`` is a mapping and a rule reads it like one. A field this + write left out is not in it, and reading it raises rather than returning + a default or a stand-in, because the value a link will deliver is not + knowable here. ``get`` keeps its own contract and answers ``None``. + """ + + def __missing__(self, key: str) -> Any: + raise _UnwrittenSibling(key) + + +class _RuleInfo: + """A pydantic ``ValidationInfo`` whose ``data`` refuses an unwritten field.""" + + __slots__ = ("_info",) + + def __init__(self, info: Any) -> None: + self._info = info + + @property + def data(self) -> Dict[str, Any]: + return _WrittenFields(self._info.data) + + def __getattr__(self, name: str) -> Any: + return getattr(self._info, name) + + +def _waiting_on_unwritten_siblings(func: Callable[..., Any]) -> Callable[..., Any]: + """Return ``func`` letting its field wait when it reads a field nobody wrote. + + Only that one field waits. Every other rule in the same write still runs, + and every other error the rule raises -- a ``KeyError`` of the rule's own + making included -- travels on and refuses the write. + """ + try: + takes_info = len(inspect.signature(func).parameters) > 1 + except (TypeError, ValueError): + takes_info = False + if not takes_info: + return func + + def waiting(value: Any, info: Any) -> Any: + try: + return func(value, _RuleInfo(info)) + except _UnwrittenSibling: + return value + + waiting.__name__ = getattr(func, "__name__", "waiting") + waiting.__qualname__ = getattr(func, "__qualname__", waiting.__name__) + return waiting + + +def _own_field_validators(model: Type[BaseModel]) -> Dict[str, Any]: + """Return ``model``'s ``mode='after'`` field validators, ready for a twin. + + Each is taken as the model bound it, so ``cls`` inside it is still the + class the user wrote it on, and ``check_fields`` is off because the twin + may be handed a subset of the fields. One reading ``info.data`` is wrapped + so that a field the write left out makes that one rule wait + (:func:`_waiting_on_unwritten_siblings`). + + Only ``mode='after'`` is taken. It judges a value the field's type has + already accepted, which is the value every later checkpoint judges too; a + ``before``, ``wrap`` or ``plain`` validator runs in place of that type and + may rewrite what it is given, and a rewritten input is refused downstream + (:func:`check_content_invariance`). + """ + rebuilt: Dict[str, Any] = {} + for name, decorator in model.__pydantic_decorators__.field_validators.items(): + info = decorator.info + if info.mode != "after": + continue + rebuilt[name] = field_validator(*info.fields, mode="after", check_fields=False)( + _waiting_on_unwritten_siblings(decorator.func) + ) + return rebuilt + + +def _write_names_of(name: str, field: Any) -> Tuple[str, ...]: + """Return every key a write may name ``field`` by: its own name and its aliases.""" + names = [name] + for alias in (field.alias, field.validation_alias): + if isinstance(alias, str): + names.append(alias) + else: + names.extend( + choice + for choice in getattr(alias, "choices", ()) + if isinstance(choice, str) + ) + return tuple(names) + + +def _fields_named_by(model: Type[BaseModel], written: FrozenSet[str]) -> FrozenSet[str]: + """Return the fields of ``model`` that ``written`` names, by name or by alias.""" + return frozenset( + name + for name, field in model.model_fields.items() + if not written.isdisjoint(_write_names_of(name, field)) + ) + + +def _rule_annotation(annotation: Any, depth: int = 0) -> Any: + """Return ``annotation`` with every model in it judged by its field rules alone. + + A nested model is replaced by :func:`_field_rules_only`, wherever it + stands -- alone, or inside a ``list[M]`` or ``dict[str, M]``. The class + the field names is still accepted where it is written: an instance was + validated when it was built. + """ + if depth > _MAX_ANNOTATION_DEPTH: + return annotation + rebuilt = _rebuild_generic(annotation, lambda arg: _rule_annotation(arg, depth + 1)) + if rebuilt is not None: + return rebuilt + if _is_model(annotation): + return Annotated[ + _field_rules_only(annotation, depth + 1), + WrapValidator(_accepting_instances_of(annotation)), + ] + return annotation + + +@functools.lru_cache(maxsize=None) +def _field_rules_only(model: Type[BaseModel], depth: int = 0) -> Type[BaseModel]: + """Return a twin of ``model`` running its field rules and not its model rules. + + A ``@model_validator`` reads the whole model, and at a write the whole + model is not there: the field it reads may be one a later write or a link + supplies. That is true of a model a field declares exactly as it is of the + model the task declares, so neither runs here. + + The twin subclasses ``model``, so a rule reading the field still gets a + value ``isinstance`` places in the class the field names and whose methods + answer. Only the annotations that name a model of their own are + redeclared, each by its own twin, so the same holds at every depth. + """ + if depth > _MAX_ANNOTATION_DEPTH: + return model + annotations: Dict[str, Any] = {} + namespace: Dict[str, Any] = { + "__annotations__": annotations, + "__module__": model.__module__, + } + for name, field in model.model_fields.items(): + rebuilt = _rule_annotation(field.annotation, depth) + if rebuilt is not field.annotation: + annotations[name] = rebuilt + namespace[name] = field + twin = cast( + Type[BaseModel], type(f"{model.__name__}__FieldRules", (model,), namespace) + ) + decorators = copy(twin.__pydantic_decorators__) + decorators.model_validators = {} + twin.__pydantic_decorators__ = decorators + twin.model_rebuild(force=True) + return twin + + +def _declares_field_rules(model: Type[BaseModel], depth: int = 0) -> bool: + """Return True when ``model`` or a model it declares carries a field rule.""" + if depth > _MAX_ANNOTATION_DEPTH: + return False + if _own_field_validators(model): + return True + return any( + _declares_field_rules(nested, depth + 1) + for field in model.model_fields.values() + for nested in _models_within(field.annotation) + ) + + +def _models_within(annotation: Any, depth: int = 0) -> List[Type[BaseModel]]: + """Return every model ``annotation`` names, however deep its containers go.""" + if depth > _MAX_ANNOTATION_DEPTH: + return [] + if _is_model(annotation): + return [annotation] + found: List[Type[BaseModel]] = [] + for arg in get_args(annotation): + found.extend(_models_within(arg, depth + 1)) + return found + + +@functools.lru_cache(maxsize=None) +def _rule_shadow( + model: Type[BaseModel], written: FrozenSet[str] +) -> Optional[Type[BaseModel]]: + """Return the shadow running ``model``'s rules over ``written``, or None if none apply. + + Only the named fields are built, and each keeps the type its model + declares: a rule reads the value its field declares, so a nested model + reaches it as an instance of the class the field names. A field left out + is absent from ``info.data`` rather than standing there as a placeholder. + + ``@model_validator``s are left out, this model's own and those of the + models its fields declare: one may read a field nobody has written. + """ + if not _declares_field_rules(model): + return None + fields = { + name: (_rule_annotation(field.annotation), _optional_field(field)) + for name, field in model.model_fields.items() + if name in written + } + if not fields: + return None + return create_model( # type: ignore[call-overload] + f"{model.__name__}__Rules", + __config__=_shadow_config(model), + __validators__=_own_field_validators(model), + **fields, + ) + + +def _rejected_wiring( + model: Type[BaseModel], exc: ValidationError, label: str +) -> TaskInputValidationError: + """Return the error reporting what ``model`` refused at ``label``.""" + return TaskInputValidationError( + f"Task '{label}' got inputs {model.__name__} rejects:\n{exc}", + task=label, + model=model, + errors=exc.errors(), + ) + + +def validate_wiring_inputs( + model: Type[BaseModel], + inputs: Dict[str, Any], + *, + label: str, + complete: bool = True, +) -> None: + """Raise unless every literal in ``inputs`` fits the field it is written to. + + ``complete`` says whether ``inputs`` is the whole call, in which case a + required field left out is refused as well. + + Types are checked first, on everything written. The model's own + ``mode='after'`` field validators then run over the fields whose value is + resolved; a field holding a socket or a task is checked for its shape + alone, because what flows through it is not known yet. A rule reading a + field nobody wrote waits for the later checkpoints, where the whole + payload is in hand -- that one rule, and not the others in the same write. + + The validated instance is discarded and ``inputs`` is passed on untouched. + That is not tidiness: pydantic strips the proxy a tagged value wears for + most field types, and a stripped value is a literal, so a task wired to + the graph's input would silently become a task holding a copy of it. + """ + from node_graph.utils import untagged_copy + + shadow = _wiring_shadow(model) if complete else _partial_wiring_shadow(model) + try: + shadow.model_validate(inputs) + except ValidationError as exc: + raise _rejected_wiring(model, exc, label) from exc + resolved = { + name: value for name, value in inputs.items() if not _awaits_a_link(value) + } + if not resolved: + return + rules = _rule_shadow(model, _fields_named_by(model, frozenset(resolved))) + if rules is None: + return + try: + rules.model_validate(untagged_copy(resolved)) + except ValidationError as exc: + raise _rejected_wiring(model, exc, label) from exc + + +def validate_task_inputs(task: Any, inputs: Dict[str, Any]) -> None: + """Check values written into ``task``'s sockets against its input model. + + Reached by ``add_task`` and by ``set_inputs``, so a value the model + refuses fails at the line that wrote it. Assigning to a socket's + ``value`` writes past this check, as it does past the socket layer's own. + A ``None`` is left out, because writing one sets nothing. + """ + model = _input_model_of_task(task) + if model is None: + return + written = {name: value for name, value in inputs.items() if value is not None} + if written: + validate_wiring_inputs( + model, written, label=getattr(task, "name", "task"), complete=False + ) + + +# -------------------------------------------------------------------------- +# Validation changes representation, never content +# -------------------------------------------------------------------------- + + +def _plain_annotation(annotation: Any, depth: int = 0) -> Any: + """Return ``annotation`` with every model replaced by its plain twin. + + Serializers written into the annotation (``Annotated[int, + PlainSerializer(...)]``) are dropped along the way. They render, and a + renderer that maps every value to the same output would hide a rewrite + from the comparison the twin exists to make. + """ + if depth > _MAX_ANNOTATION_DEPTH: + return annotation + if get_origin(annotation) is Annotated: + args = get_args(annotation) + inner = _plain_annotation(args[0], depth + 1) + kept = [item for item in args[1:] if not isinstance(item, _SERIALIZER_METADATA)] + return Annotated[tuple([inner, *kept])] if kept else inner + rebuilt = _rebuild_generic( + annotation, lambda arg: _plain_annotation(arg, depth + 1) + ) + if rebuilt is not None: + return rebuilt + if _is_model(annotation): + return _plain_twin(annotation, depth + 1) + return annotation + + +def _comparable_field(field: Any) -> Any: + """Return ``field`` with what only renders it taken off, so it is comparable. + + Two things are taken off. ``Field(exclude=True)`` keeps a field out of a + dump, which would keep it out of the comparison too. A serializer written + into the annotation (``Annotated[int, PlainSerializer(...)]``) lands in + the field's metadata rather than in its type, and one that maps every + value to the same output would hide a rewrite behind it. Both say how a + value is rendered, not what was written. + """ + metadata = list(getattr(field, "metadata", ()) or ()) + kept = [item for item in metadata if not isinstance(item, _SERIALIZER_METADATA)] + if getattr(field, "exclude", None) is None and len(kept) == len(metadata): + return field + field = copy(field) + field.exclude = None + field.metadata = kept + return field + + +@functools.lru_cache(maxsize=None) +def _plain_twin(model: Type[BaseModel], depth: int = 0) -> Type[BaseModel]: + """Return ``model``'s fields with its validators and serializers left out. + + Built with no base class, so what survives is the field types and their + constraints, and ``model_config``, which decides how a value is coerced. + Coercion is therefore identical to the model's and every rule the user + wrote is absent, which is what makes the twin a reference for what the + input said before any rule ran. + + ``depth`` counts the models the rebuild has walked through, so a model + that names itself terminates. + """ + if depth > _MAX_ANNOTATION_DEPTH: + return model + fields = { + name: (_plain_annotation(field.annotation, depth), _comparable_field(field)) + for name, field in model.model_fields.items() + } + return create_model( # type: ignore[call-overload] + f"{model.__name__}__Content", + __config__=_shadow_config(model), + **fields, + ) + + +def _plain_values(value: Any) -> Any: + """Return ``value`` with model instances turned into plain dicts. + + Field values are read as attributes rather than dumped, so a + ``field_serializer`` -- which renders, and so is representation -- does not + reach the comparison. + """ + if isinstance(value, BaseModel): + return { + name: _plain_values(getattr(value, name)) + for name in type(value).model_fields + } + if isinstance(value, dict): + return {key: _plain_values(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return type(value)(_plain_values(item) for item in value) + return value + + +def _content_of( + model: Type[BaseModel], values: Any +) -> Tuple[Optional[Dict[str, Any]], Optional[str]]: + """Return the content ``model``'s twin reads in ``values``, field by field. + + A field is rendered as JSON, which is the form two values can be compared + in whatever they are spelled as. A field pydantic cannot render -- an + engine's data node or any object under ``Any`` -- keeps the object itself + instead, so it is compared as it stands rather than aborting the check + with a serialization error. + """ + twin = _plain_twin(model) + try: + instance = twin.model_validate(_plain_values(values)) + except ValidationError as exc: + return None, str(exc) + content: Dict[str, Any] = {} + for name in type(instance).model_fields: + try: + content[name] = instance.model_dump( + mode="json", include={name}, warnings=False + )[name] + except (PydanticSerializationError, KeyError): + content[name] = getattr(instance, name) + return content, None + + +def _same_content(before: Any, after: Any) -> bool: + """Return True when two readings of a field say the same thing. + + A field rendered as JSON compares as JSON. One kept as the object it is -- + an engine's data node, anything under ``Any`` -- may answer an inequality + with something other than a bool: a numpy array answers with an array of + them, which no ``if`` can read. Such a value counts as unchanged when it + is the same object, when a whole-array comparison says so, or when the two + render identically. + """ + if before is after: + return True + try: + return not bool(before != after) + except (TypeError, ValueError): + pass + numpy = sys.modules.get("numpy") + if ( + numpy is not None + and isinstance(before, numpy.ndarray) + and isinstance(after, numpy.ndarray) + ): + return bool(numpy.array_equal(before, after)) + return repr(before) == repr(after) + + +class ContentSnapshot(NamedTuple): + """What a model's inputs said, read before any of its validators ran.""" + + content: Optional[Dict[str, Any]] + error: Optional[str] + names: Tuple[str, ...] + + +def content_snapshot(model: Type[BaseModel], given: Dict[str, Any]) -> ContentSnapshot: + """Return what ``given`` says, read before ``model`` is allowed to run. + + Taken first and kept, because a validator may rewrite the very values it + is judged against: one that appends to a list in place, or reaches into a + nested model instance, leaves nothing behind to compare with. Rendering + the values to JSON here is what makes the snapshot independent of what + happens to them afterwards. + """ + content, error = _content_of(model, given) + return ContentSnapshot(content, error, tuple(given)) + + +def check_content_invariance( + model: Type[BaseModel], + before_snapshot: ContentSnapshot, + validated: BaseModel, + *, + label: str, +) -> None: + """Raise if validating changed what any supplied field says. + + ``before_snapshot`` comes from :func:`content_snapshot`, taken before the + model ran. Coercion is free to change how a value is spelled -- ``'60'`` + may become a ``Decimal`` and a list a tuple -- because both spellings + carry the same content. Deriving a value is not: the body would then run + on a value that never reached storage, so provenance would record one + input and the body would have seen another. A rule that leaves an + already-resolved value alone therefore passes, and one that rewrites it is + refused. + + Only the fields the caller supplied are compared, so a default filling a + field the caller omitted is not a change. Validators written into an + annotation (``Annotated[int, AfterValidator(...)]``) are part of the type + and run on both sides of the comparison, so they must be idempotent: one + that doubles its input is refused, because the second run doubles it + again. + + A field the model has no JSON form for is compared as the object it is, + which catches a rule that replaces it and not one that reaches inside it. + """ + before, before_error, names = before_snapshot + after, after_error = _content_of(model, _plain_values(validated)) + if before_error is not None: + raise ModelDerivedValueError( + f"'{label}' cannot check what {model.__name__} did to its inputs: " + f"the values it was given do not fit the fields they were written to.\n{before_error}\n" + "How to fix: a validator that rewrites a value before it is checked has no reference to " + "compare against; widen the field's type to accept what the caller writes, or derive the " + "value before the task is built." + ) + if after_error is not None or before is None or after is None: + raise ModelDerivedValueError( + f"'{label}' got values from {model.__name__} that its own fields refuse.\n{after_error}" + ) + for name in names: + if name not in before or name not in after: + continue + if not _same_content(before[name], after[name]): + raise ModelDerivedValueError( + f"'{label}': validating input '{name}' through {model.__name__} changed it from " + f"{before[name]!r} to {after[name]!r}.\n" + "How to fix: validation may change how a value is spelled but not what it says. " + "Derive the value in the outermost model, before the graph is built, and pass the " + "derived value in." + ) + + +# -------------------------------------------------------------------------- +# Checkpoint B -- a graph's resolved inputs +# -------------------------------------------------------------------------- + + +def _under_the_original_tags(written: Any, rebuilt: Any) -> Any: + """Return ``rebuilt`` wearing the tags ``written`` wore, leaf by leaf. + + A tag is what a graph body turns into a link, so a value rebuilt without + one would leave the body holding a copy of the graph's input. A mapping is + walked rather than replaced, because its members carry tags of their own. + """ + from node_graph.socket import TaggedValue + from node_graph.utils.struct_utils import retagged + + if isinstance(written, TaggedValue): + return retagged(_under_the_original_tags(written.__wrapped__, rebuilt), written) + if isinstance(written, dict): + if isinstance(rebuilt, BaseModel): + return { + key: _under_the_original_tags(value, getattr(rebuilt, key, value)) + for key, value in written.items() + } + if isinstance(rebuilt, dict): + return { + key: _under_the_original_tags(value, rebuilt.get(key, value)) + for key, value in written.items() + } + return rebuilt + + +def validate_graph_inputs( + model: Type[BaseModel], + inputs: Dict[str, Any], + *, + label: str, +) -> Dict[str, Any]: + """Return a graph's inputs in the types ``model`` declares, or raise. + + The graph's inputs are values by the time its body runs, so the real model + runs here, cross-field rules included, and the body is handed what the + model made of each field: an ``Enum`` field arrives as the member, a + ``Decimal`` field as a ``Decimal``. Every value comes back under the tag + the caller wrote it with, and under that tag's uuid, because the body + turns a tag into a link and provenance identifies a value by the uuid. + + Every field of the model is handed over, whether or not a value arrived + for it: a scalar nobody wrote arrives at its default and a nested model or + a mapping arrives as the members that were written, which is nothing at + all when the field itself was left out. + + ``inputs`` arrive as the body will see them: the engine's read edge has + already taken off the node each leaf was stored in. + """ + from node_graph.utils import untagged_copy + + given = untagged_copy(inputs) + before = content_snapshot(model, given) + try: + validated = _by_field_name(model).model_validate(given) + except ValidationError as exc: + raise TaskInputValidationError( + f"Graph '{label}' got inputs {model.__name__} rejects:\n{exc}", + task=label, + model=model, + errors=exc.errors(), + ) from exc + check_content_invariance(model, before, validated, label=label) + handed = body_inputs(model, validated) + return { + name: _under_the_original_tags(inputs[name], value) if name in inputs else value + for name, value in handed.items() + } + + +# -------------------------------------------------------------------------- +# Checkpoint C -- the run edge +# -------------------------------------------------------------------------- + + +def _dump_model_instance(instance: BaseModel) -> Dict[str, Any]: + """Return the JSON-safe form of every field of ``instance``. + + A field the model has no JSON form for keeps the object it holds, exactly + as :func:`dump_model_field` does on the way in, so one such output does + not cost the whole return value its rendering. + """ + try: + return cast(Dict[str, Any], instance.model_dump(mode="json", warnings=False)) + except PydanticSerializationError: + pass + dumped: Dict[str, Any] = {} + for name in type(instance).model_fields: + try: + dumped[name] = instance.model_dump( + mode="json", include={name}, warnings=False + )[name] + except (PydanticSerializationError, KeyError): + dumped[name] = getattr(instance, name) + return dumped + + +def _refuse_undeclared_outputs(model: Type[BaseModel], result: Any, label: str) -> None: + """Raise when a returned mapping carries a key ``model`` does not declare. + + The model declares the output sockets, so a key it does not name has + nowhere to be written and would be dropped between the body and the + task's outputs. + """ + if not isinstance(result, dict): + return + declared = set() + for name, field in model.model_fields.items(): + declared.add(name) + for alias in (field.alias, field.validation_alias, field.serialization_alias): + if isinstance(alias, str): + declared.add(alias) + undeclared = [key for key in result if key not in declared] + if undeclared: + raise TaskOutputValidationError( + f"Task '{label}' returned {', '.join(repr(key) for key in undeclared)}, " + f"which {model.__name__} does not declare.\n" + f"How to fix: add the field(s) to {model.__name__}, or drop them from what " + f"the body returns." + ) + + +#: Suffix under which a module answers for a task's validated wrapper. +_VALIDATED_SUFFIX = "__node_graph_validated" + + +def _bind_validated_callable( + wrapper: Callable[..., Any], func: Callable[..., Any] +) -> None: + """Bind ``wrapper`` in ``func``'s module under the name its executor records. + + A task records its executor as ``.`` and resolves + it by importing that path again, so the name recorded has to reach the + wrapper rather than ``func``, which carries no model and enforces + nothing. ``functools.wraps`` gives the wrapper ``func``'s own name, so + the wrapper is bound here under a name of its own and that is the name + the executor records. + + ``__name__`` is left as it was, because it is what an engine labels a + process with, and so is ``__qualname__``: it decides that a callable + defined in ``__main__`` or nested in another scope travels as a pickle + instead, and a pickle carries the models with it. + """ + from node_graph.executor import BOUND_NAME_ATTR + + module_name = getattr(func, "__module__", None) + qualname = getattr(func, "__qualname__", "") or "" + if not module_name or module_name == "__main__" or "." in qualname: + return + base = f"{getattr(func, '__name__', 'task')}{_VALIDATED_SUFFIX}" + module = sys.modules.get(module_name) + if module is None: + # Importing the module runs this decoration again and binds the name. + setattr(wrapper, BOUND_NAME_ATTR, base) + return + name, index = base, 1 + while getattr(module, name, wrapper) is not wrapper: + index += 1 + name = f"{base}_{index}" + setattr(wrapper, BOUND_NAME_ATTR, name) + setattr(module, name, wrapper) + + +def _written_members(instance: BaseModel) -> Dict[str, Any]: + """Return the members written into ``instance``, in the types it made of them. + + A member nobody wrote is left out at every depth. The values are the + model's own -- a ``Decimal`` field is a ``Decimal`` here, not the string + its serializer stores -- because this is what the body is called with, + not what storage keeps. + """ + written = instance.__pydantic_fields_set__ + members: Dict[str, Any] = {} + for name in type(instance).model_fields: + if name not in written: + continue + value = getattr(instance, name) + members[name] = ( + _written_members(value) if isinstance(value, BaseModel) else value + ) + return members + + +def body_inputs(model: Type[BaseModel], validated: BaseModel) -> Dict[str, Any]: + """Return what the body is called with, field by field. + + A field declaring a nested model is handed the members that were written + and no others, so a body reading ``system`` reads the keys its caller + named, and a ``dict[str, T]`` field is handed each of its items on those + same terms -- a mapping is one namespace per key, so the rule holds one + level further down. Every other field is handed what the model made of + it, default included: a field the caller left out is one the model + answers for. + """ + inputs: Dict[str, Any] = {} + for name, value in dict(validated).items(): + field = model.model_fields.get(name) + annotation = field.annotation if field is not None else None + nested = _model_of_type(annotation) if annotation is not None else None + if nested is not None and isinstance(value, BaseModel): + inputs[name] = _written_members(value) + continue + item_type = _container_item_type(annotation) if annotation is not None else None + item_model = _model_of_type(item_type) if item_type is not None else None + if item_model is not None and isinstance(value, dict): + inputs[name] = { + key: _written_members(item) if isinstance(item, BaseModel) else item + for key, item in value.items() + } + continue + inputs[name] = value + return inputs + + +def rebind_executor_callable( + outer: Callable[..., Any], inner: Callable[..., Any] +) -> None: + """Move the name ``inner`` is bound under onto ``outer``. + + A callable wrapping a validated callable -- an engine's own decorator over + it -- is what a task's executor has to resolve to, while the validated + callable goes on enforcing the models from inside. The binding moves, so + the name the executor records reaches the outer callable. + """ + from node_graph.executor import BOUND_NAME_ATTR + + name = getattr(inner, BOUND_NAME_ATTR, None) + if name is None: + return + module = sys.modules.get(getattr(inner, "__module__", "") or "") + if module is not None and getattr(module, name, None) is inner: + setattr(module, name, outer) + setattr(outer, BOUND_NAME_ATTR, name) + + +def validated_callable( + func: Callable[..., Any], + *, + input_model: Optional[Type[BaseModel]] = None, + output_model: Optional[Type[BaseModel]] = None, + label: Optional[str] = None, +) -> Callable[..., Any]: + """Return ``func`` wrapped so its models are enforced around every call. + + The wrapper is what the task's executor points at, so validation happens + wherever the body runs: inputs are validated once the engine has assembled + and deserialized them, outputs once the body has returned. The body sees + the values the input model produced; the caller sees the JSON-safe form of + what the output model accepted. + """ + label = label or getattr(func, "__name__", "task") + signature = inspect.signature(func) + + @functools.wraps(func) + def call(*args: Any, **kwargs: Any) -> Any: + if args: + # An engine that calls the body positionally -- AiiDA's process + # functions do -- addresses the same fields by name here. + kwargs = dict(signature.bind_partial(*args, **kwargs).arguments) + if input_model is not None: + before = content_snapshot(input_model, kwargs) + try: + validated = _by_field_name(input_model).model_validate(kwargs) + except ValidationError as exc: + raise TaskInputValidationError( + f"Task '{label}' got inputs {input_model.__name__} rejects:\n{exc}", + task=label, + model=input_model, + errors=exc.errors(), + ) from exc + check_content_invariance(input_model, before, validated, label=label) + kwargs = body_inputs(input_model, validated) + result = func(**kwargs) + if output_model is not None: + _refuse_undeclared_outputs(output_model, result, label) + try: + accepted = _by_field_name(output_model).model_validate(result) + except ValidationError as exc: + raise TaskOutputValidationError( + f"Task '{label}' returned outputs {output_model.__name__} rejects:\n{exc}" + ) from exc + result = _dump_model_instance(accepted) + return result + + setattr(call, INPUT_MODEL_ATTR, input_model) + setattr(call, OUTPUT_MODEL_ATTR, output_model) + _bind_validated_callable(call, func) + return call + + +# -------------------------------------------------------------------------- +# Storage -- the model renders its own leaves +# -------------------------------------------------------------------------- + + +def dump_model_field(model: Type[BaseModel], name: str, value: Any) -> Any: + """Return the JSON-safe form ``model`` gives ``name``'s ``value``. + + The value is placed in an unvalidated instance and dumped through the + model, so a ``field_serializer`` declared on ``model`` decides the stored + form. Only ``name`` is dumped, so the other fields need not be present. + + A value the model has no JSON form for -- an engine's data node, an object + under ``Any`` -- and a field ``Field(exclude=True)`` keeps out of a dump + are returned as they are, for the engine's own serialization to store or + to refuse in its own words. + """ + holder = model.model_construct(**{name: value}) + try: + return holder.model_dump(mode="json", include={name}, warnings=False)[name] + except (PydanticSerializationError, KeyError): + return value + + +def dump_model_item(model: Type[BaseModel], name: str, key: str, value: Any) -> Any: + """Return the JSON-safe form ``model`` gives one entry of its ``name`` mapping.""" + return dump_model_field(model, name, {key: value})[key] + + +def input_model_of_callable(func: Any) -> Optional[Type[BaseModel]]: + """Return the input model ``func`` enforces, if it enforces one.""" + return getattr(func, INPUT_MODEL_ATTR, None) + + +def output_model_of_callable(func: Any) -> Optional[Type[BaseModel]]: + """Return the output model ``func`` enforces, if it enforces one.""" + return getattr(func, OUTPUT_MODEL_ATTR, None) + + +def _input_model_of_task(task: Any) -> Optional[Type[BaseModel]]: + """Return the input model of ``task``, resolving its executor once. + + A task whose executor names a callable that cannot be reached from here -- + a module that is not importable, a spec with no executor at all -- has no + contract to apply, and the generic path runs as it did before models. + + Any other failure is a model this process cannot read, and a model it + cannot read is one it cannot enforce: that is a + :class:`ModelContractError`, not a task that quietly loses its checks. + """ + cached = getattr(task, _TASK_MODEL_CACHE, _UNRESOLVED) + if cached is not _UNRESOLVED: + return cast(Optional[Type[BaseModel]], cached) + model = None + try: + executor = task.spec.executor + callable_obj = executor.callable if executor is not None else None + callable_obj = getattr(callable_obj, "_callable", callable_obj) + model = input_model_of_callable(callable_obj) + except (ImportError, AttributeError): + model = None + except Exception as exc: + raise ModelContractError( + f"Task '{getattr(task, 'name', '?')}' carries a callable this process " + f"cannot rebuild, so no input model can be read from it: " + f"{type(exc).__name__}: {exc}\n" + "How to fix: move the task's function and the models it declares " + "into an importable module." + ) from exc + try: + setattr(task, _TASK_MODEL_CACHE, model) + except AttributeError: + pass + return model + + +def _dumper_for_path( + model: Type[BaseModel], path: List[str] +) -> Optional[Callable[[Any], Any]]: + """Return the function rendering the value at ``path`` under ``model``. + + ``path`` walks the socket names below the task's inputs. A name is a field + of the model reached so far, except directly after a ``dict[str, T]`` + field, where it is one of the mapping's keys and the walk continues in + ``T``. The walk ends at the model that declares the addressed value, which + is the model whose serializers render it. + """ + current: Optional[Type[BaseModel]] = model + index = 0 + while current is not None and index < len(path): + owner, name = current, path[index] + field = owner.model_fields.get(name) + if field is None: + return None + if index == len(path) - 1: + return functools.partial(dump_model_field, owner, name) + + item_type = _container_item_type(field.annotation) + if item_type is None: + current = _model_of_type(field.annotation) + index += 1 + continue + if index + 1 == len(path) - 1: + return functools.partial(dump_model_item, owner, name, path[index + 1]) + current = _model_of_type(item_type) + index += 2 + return None + + +def model_dumper_for_socket(socket: Any) -> Optional[Callable[[Any], Any]]: + """Return the function rendering a socket's stored value, if a model owns it. + + A socket is owned when its path resolves to a declared field, however deep + the models and typed mappings it passes through. A socket the walk cannot + place -- under a leaf ``dict``, say -- falls through to the generic + serialization path, which will refuse a value it has no serializer for. + """ + full_name = getattr(socket, "_full_name", None) + task = getattr(socket, "_task", None) + if not full_name or task is None: + return None + parts = full_name.split(".") + if len(parts) < 2 or parts[0] != "inputs": + return None + model = _input_model_of_task(task) + if model is None: + return None + return _dumper_for_path(model, parts[1:]) + + +# -------------------------------------------------------------------------- +# Decoration +# -------------------------------------------------------------------------- + + +def apply_models( + obj: Callable[..., Any], + inputs: Any, + outputs: Any, + input_model: Optional[Type[BaseModel]], + output_model: Optional[Type[BaseModel]], + *, + is_graph: bool = False, + api: Any = None, +) -> Tuple[Any, Any, Callable[..., Any]]: + """Return ``(inputs, outputs, callable)`` for a task built from models. + + Without either model this is the identity, so a task declared the usual + way takes exactly the path it took before. A graph task keeps its own + function as the callable -- the engine reinserts a recursion handle into + that function's globals, which a wrapper does not carry -- and is + validated where its inputs are resolved instead. + """ + if input_model is None and output_model is None: + return inputs, outputs, obj + + if not inspect.isfunction(obj) or getattr(obj, "node_class", False): + raise ModelContractError( + "input_model/output_model apply to a plain Python function task.\n" + f"How to fix: drop them from {obj!r}, or wrap the process in a function task." + ) + if input_model is not None and inputs is not None: + raise ModelContractError( + "'inputs' and 'input_model' both declare the input sockets.\n" + "How to fix: keep input_model and drop inputs." + ) + if output_model is not None and outputs is not None: + raise ModelContractError( + "'outputs' and 'output_model' both declare the output sockets.\n" + "How to fix: keep output_model and drop outputs." + ) + if is_graph and output_model is not None: + raise ModelContractError( + "output_model is not supported on @task.graph.\n" + "Reason: a graph body returns socket references, which stand for values that do not exist " + "yet, so nothing can be validated against the model.\n" + "How to fix: put output_model on the function tasks whose outputs the graph returns." + ) + + if input_model is not None: + check_signature_against_model(obj, input_model, role="input") + inputs = spec_from_model(input_model, api) + if output_model is not None: + check_signature_against_model(obj, output_model, role="output") + outputs = spec_from_model(output_model, api) + + if is_graph: + # The graph body itself is the executor, so the models ride on it. + setattr(obj, INPUT_MODEL_ATTR, input_model) + setattr(obj, OUTPUT_MODEL_ATTR, output_model) + return inputs, outputs, obj + + wrapped = validated_callable( + obj, input_model=input_model, output_model=output_model + ) + return inputs, outputs, wrapped diff --git a/src/node_graph/link.py b/src/node_graph/link.py index 2e2d535..062f890 100644 --- a/src/node_graph/link.py +++ b/src/node_graph/link.py @@ -85,6 +85,144 @@ def _annotated_type(self, sock: "Socket") -> str | None: return self._format_union(union) return self._annotated_py_type(sock) + def _socket_extras(self, sock: "Socket") -> dict: + return getattr(getattr(sock, "_metadata", None), "extras", {}) or {} + + def _allowed_values(self, sock: "Socket") -> list | None: + """Return the values this socket may carry, or None if unrestricted. + + A ``Literal`` socket lists them outright; a plain ``Enum`` socket + carries all of its members' values. + """ + extras = self._socket_extras(sock) + allowed = extras.get("allowed_values") + if isinstance(allowed, list): + return allowed + info = extras.get("structured_type") + if isinstance(info, dict) and info.get("kind") == "enum": + from node_graph.utils.struct_utils import import_structured_type + + try: + return [member.value for member in import_structured_type(info["path"])] + except Exception: + return None + return None + + def _literal_base(self, sock: "Socket") -> str | None: + base = self._socket_extras(sock).get("literal_base") + return str(base).lower() if base else None + + def check_allowed_values(self) -> bool | None: + """Decide a link where either end restricts the values it carries. + + Returns True to accept the link, None to leave the decision to the + remaining checks, and raises when the source can carry a value the + target forbids. Reached only for a source whose type is known: an + ``Any`` source is accepted earlier on type alone, which is what lets an + untyped parent pass a value on, and ``check_static_source_value`` reads + whatever value it already holds. + """ + from node_graph.utils.struct_utils import value_is_allowed + + from_allowed = self._allowed_values(self.from_socket) + to_allowed = self._allowed_values(self.to_socket) + if from_allowed is None and to_allowed is None: + return None + + if from_allowed is not None and to_allowed is not None: + offending = [ + value + for value in from_allowed + if not value_is_allowed(value, to_allowed) + ] + if not offending: + return True + self._raise_allowed_values_mismatch( + [repr(value) for value in offending], to_allowed + ) + + if to_allowed is not None: + # The source is unrestricted, so it can carry a forbidden value. + self._raise_allowed_values_mismatch(None, to_allowed) + + # Only the source restricts: widening into its own base type is safe. + base = self._literal_base(self.from_socket) + if base is not None and self._lower_id(self.to_socket) == base: + return True + return None + + def check_static_source_value(self) -> None: + """Reject a link whose source already holds a value the target forbids. + + A source socket declared ``Any`` passes every type check, so a value + already sitting on it is read here, the last point before the run. + Everything else is decided at run time, when the value exists: a source + fed by an upstream task, a source assigned after this link was made, + and a value bound for a sub-graph one hop further out, whose restricted + socket only appears as that sub-graph expands. + """ + from node_graph.socket import TaggedValue + + from node_graph.utils.struct_utils import ( + canonical_socket_value, + socket_loc, + socket_subject, + ) + + extras = self._socket_extras(self.to_socket) + structured_type = extras.get("structured_type") or {} + if "allowed_values" not in extras and structured_type.get("kind") != "enum": + return + if self._is_namespace(self.from_socket): + # A namespace carries a mapping; the shape check below reports it. + return + if any(link.to_socket is self.from_socket for link in self.from_socket._links): + # Fed by an upstream task: the value only exists at run time. + return + value = getattr(self.from_socket, "_value", None) + if value is None: + return + if isinstance(value, TaggedValue): + tagged_socket = value._socket + # ``is`` only: comparing sockets with ``==`` builds an operator task. + if tagged_socket is not None and tagged_socket is not self.from_socket: + return + value = value.__wrapped__ + canonical_socket_value( + value, + structured_type=structured_type or None, + allowed=extras.get("allowed_values"), + subject=socket_subject(self.to_socket._full_name_with_task), + loc=socket_loc(self.to_socket._full_name), + ) + + def _raise_allowed_values_mismatch( + self, offending: list[str] | None, allowed: list + ) -> None: + src = f"{self.from_task.name}.{self.from_socket._scoped_name}" + dst = f"{self.to_task.name}.{self.to_socket._scoped_name}" + rendered = ", ".join(repr(value) for value in allowed) + detail = ( + f" Source can carry {', '.join(offending)}, which {dst} forbids." + if offending + else " Source is not restricted to those values." + ) + raise TypeError( + "\n".join( + [ + "Socket value range mismatch:", + f" {src} [{self._format_socket_id(self.from_socket)}] -> {dst} " + f"[{self._format_socket_id(self.to_socket)}] is not allowed.", + f" {dst} accepts only {rendered}.", + detail, + "", + "Suggestions:", + " • Widen the target annotation to include the source's values", + " • Or insert a task that narrows the value before the link", + ] + ) + ) + def _namespace_item(self, sock: "Socket") -> dict | None: extras = getattr(getattr(sock, "_metadata", None), "extras", {}) or {} item = extras.get("item") @@ -169,6 +307,8 @@ def check_socket_match(self) -> None: from_id = self._lower_id(self.from_socket) to_id = self._lower_id(self.to_socket) + self.check_static_source_value() + # "any" accepts anything if self._is_any(self.from_socket) or self._is_any(self.to_socket): return @@ -254,6 +394,9 @@ def check_socket_match(self) -> None: ): return + if self.check_allowed_values(): + return + if self._is_annotated(self.from_socket) and self._is_annotated(self.to_socket): from_type = self._annotated_py_type(self.from_socket) to_type = self._annotated_py_type(self.to_socket) diff --git a/src/node_graph/socket.py b/src/node_graph/socket.py index 6ad9f5b..4168117 100644 --- a/src/node_graph/socket.py +++ b/src/node_graph/socket.py @@ -708,6 +708,8 @@ def _set_socket_value(self, value: Any, *, value_source: str = "link") -> None: elif is_input and value_source != "property": self._update_updatable_meta({"value_source": "link"}) + value = self._canonical_value(value) + graph = getattr(self._task, "graph", None) if graph is not None: policy = getattr(graph, "serialization_policy", "off") @@ -722,6 +724,52 @@ def _set_socket_value(self, value: Any, *, value_source: str = "link") -> None: f"Socket '{self._name}' has no property to set a value." ) + def _canonical_value(self, value: Any) -> Any: + """Return the value this socket accepts, raising when it accepts none. + + Input sockets built from an ``Enum`` or a ``Literal`` carry the members + they admit; every other socket returns ``value`` unchanged. What is + stored for an ``Enum`` is the member's value, the one form that also + comes back out of storage, so a socket reads the same whichever side of + a process boundary filled it; ``coerce_inputs_from_spec`` rebuilds the + member for the task body. ``None`` passes through, so an optional socket + can still be cleared. Output sockets are left alone: a task reports what + it computed. + """ + from node_graph.utils.struct_utils import ( + canonical_socket_value, + literal_value, + retagged, + socket_loc, + socket_subject, + ) + + if self._full_name.split(".")[0] != "inputs": + return value + extras = self._metadata.extras or {} + structured_type = extras.get("structured_type") + allowed = extras.get("allowed_values") + if allowed is None and ( + structured_type is None or structured_type.get("kind") != "enum" + ): + return value + if value is None: + return value + + raw = value.__wrapped__ if isinstance(value, TaggedValue) else value + canonical = literal_value( + canonical_socket_value( + raw, + structured_type=structured_type, + allowed=allowed, + subject=socket_subject(self._full_name_with_task), + loc=socket_loc(self._full_name), + ) + ) + if canonical is raw: + return value + return retagged(canonical, value) + def _serialize_value(self, store: bool = False) -> Any: """Serialize the socket value unless it's metadata (stored as raw).""" value = resolve_tagged_values(self._value) @@ -850,8 +898,19 @@ def _raise_namespace_assignment_error( incoming_desc: str, reason: str, fixes: list[str], + error_type: str, + key: str | None = None, ) -> None: - """Raise a ValueError guiding users when setting/linking values into a namespace.""" + """Raise the refusal of a value written or linked into a namespace. + + ``error_type`` names the rule that refused and ``key``, where the refusal + is about one name below the namespace, is that name: together with the + namespace's own path they address the refused value the way an input + model's errors address it. + """ + from node_graph.errors import SocketValueError + from node_graph.utils.struct_utils import socket_loc + where = getattr(target_ns, "_full_name_with_task", "") msg = [ f"Invalid assignment into namespace socket: {where}", @@ -867,7 +926,11 @@ def _raise_namespace_assignment_error( "How to fix:", *[f" • {line}" for line in fixes], ] - raise ValueError("\n".join(msg)) + raise SocketValueError( + "\n".join(msg), + loc=socket_loc(getattr(target_ns, "_full_name", None), *([key] if key else [])), + error_type=error_type, + ) class TaskSocketNamespace(BaseSocket, OperatorSocketMixin): @@ -1214,6 +1277,8 @@ def _get_or_create_child(self, key: str, val: Any) -> object: "Remove the builtin key from the assignment; or", "link to an explicit leaf socket instead of the namespace.", ], + error_type="reserved_key", + key=key, ) if not self._metadata.dynamic: _raise_namespace_assignment_error( @@ -1224,6 +1289,8 @@ def _get_or_create_child(self, key: str, val: Any) -> object: "Add the field to the namespace spec; or", "make the namespace dynamic if it should grow automatically.", ], + error_type="extra_forbidden", + key=key, ) self._append_dynamic_child(key, val) return self._sockets[key] @@ -1241,6 +1308,8 @@ def _assign_key_value(self, key: str, val: Any, *, value_source: str) -> None: "Remove the builtin key from the assignment; or", "link to an explicit leaf socket instead of the namespace.", ], + error_type="reserved_key", + key=head, ) if not self._metadata.dynamic: _raise_namespace_assignment_error( @@ -1252,6 +1321,8 @@ def _assign_key_value(self, key: str, val: Any, *, value_source: str) -> None: "mark this namespace dynamic if it must accept arbitrary keys;", "or correct the key path being assigned.", ], + error_type="extra_forbidden", + key=head, ) self._new( self._SocketPool["namespace"], @@ -1271,6 +1342,8 @@ def _assign_key_value(self, key: str, val: Any, *, value_source: str) -> None: "Use a namespace socket for hierarchical data; update your SocketSpec accordingly; or", "flatten your assignment to target a leaf socket directly.", ], + error_type="not_a_namespace", + key=head, ) child._set_socket_value({tail: val}, value_source=value_source) return @@ -1297,6 +1370,7 @@ def _assign_key_value(self, key: str, val: Any, *, value_source: str) -> None: "Link a specific leaf socket instead of the namespace; or", "make the namespace dynamic if it should accept multiple links.", ], + error_type="namespace_link_forbidden", ) if isinstance(val, (dict, TaskSocketNamespace)): target._set_socket_value(val, value_source=value_source) @@ -1308,6 +1382,7 @@ def _assign_key_value(self, key: str, val: Any, *, value_source: str) -> None: fixes=[ "Provide a dict like {'x': 1, 'y': 2} that matches the namespace shape.", ], + error_type="namespace_expects_mapping", ) else: target._set_socket_value(val, value_source=value_source) @@ -1355,6 +1430,7 @@ def _set_socket_value( " add_multiply(data=data)", "", ], + error_type="namespace_tagged_value", ) if not isinstance(value, dict): @@ -1366,6 +1442,7 @@ def _set_socket_value( "Provide a dict with keys matching the namespace fields;", "or link another socket/namespace; or declare the graph input as a namespace.", ], + error_type="namespace_expects_mapping", ) for key, val in value.items(): diff --git a/src/node_graph/socket_spec.py b/src/node_graph/socket_spec.py index 2a1d479..d45818a 100644 --- a/src/node_graph/socket_spec.py +++ b/src/node_graph/socket_spec.py @@ -11,9 +11,18 @@ ) import inspect from copy import deepcopy +from enum import Enum from node_graph.orm.mapping import type_mapping as DEFAULT_TM from node_graph.socket_meta import CallRole, SocketMeta, merge_meta -from node_graph.utils.struct_utils import structured_type_info +from node_graph.utils.struct_utils import ( + JSON_SAFE_LITERALS, + canonical_socket_value, + is_enum_type, + literal_value, + structured_type_info, + structured_type_path, + value_is_allowed, +) from .socket import TaskSocketNamespace import ast import textwrap @@ -24,7 +33,7 @@ MISSING as _DC_MISSING, ) -from typing import Annotated, get_args, get_origin, get_type_hints +from typing import Annotated, Literal, get_args, get_origin, get_type_hints try: from typing import Unpack @@ -68,6 +77,17 @@ def _is_union_origin(origin: Any) -> bool: return (origin is Union) or (_UNION_TYPE is not None and origin is _UNION_TYPE) +def _literal_arg_name(arg: Any) -> str: + """Render one ``Literal`` argument. + + An ``Enum`` member is written out with its class's import path, so two + same-named enums from different modules do not render alike. + """ + if isinstance(arg, Enum): + return f"{structured_type_path(type(arg))}.{arg.name}" + return repr(arg) + + def _is_annotated_type(tp: Any) -> bool: """True if tp is an Annotated[...] wrapper across versions.""" return get_origin(tp) is Annotated @@ -334,6 +354,35 @@ class SocketSpec: fields: Dict[str, "SocketSpec"] = field(default_factory=dict) meta: SocketMeta = field(default_factory=SocketMeta) + def __post_init__(self) -> None: + """Hold the default a socket assignment would have accepted. + + A default reaches ``TaskProperty.value`` without passing the socket's + own setter, so it is decided here instead: a default an ``Enum`` or + ``Literal`` socket forbids raises where the spec is built, and one it + admits is stored as the bare value the socket would have kept. + """ + if isinstance(self.default, type(MISSING)) or self.default is None: + return + extras = self.meta.extras or {} + structured_type = extras.get("structured_type") + allowed = extras.get("allowed_values") + if allowed is None and ( + structured_type is None or structured_type.get("kind") != "enum" + ): + return + canonical = literal_value( + canonical_socket_value( + self.default, + structured_type=structured_type, + allowed=allowed, + subject=f"the default of a socket annotated " + f"{extras.get('py_type', self.identifier)}", + ) + ) + if canonical is not self.default: + object.__setattr__(self, "default", canonical) + @property def dynamic(self) -> bool: return bool(self.meta.dynamic) @@ -408,7 +457,7 @@ def __getattr__(self, name: str) -> "SocketView": raise AttributeError("'.item' only valid on dynamic namespace specs") if spec.fields and name in spec.fields: return SocketView(spec.fields[name]) - raise AttributeError(f"'{name}' not found in namespace spec") + raise AttributeError(f"namespace spec has no field '{name}'") def __getitem__(self, name: str) -> "SocketView": return self.__getattr__(name) @@ -551,6 +600,9 @@ def _map_identifier(cls, tp: Any) -> str: @staticmethod def _py_type_name(tp: Any) -> str: + if get_origin(tp) is Literal: + args = ", ".join(_literal_arg_name(arg) for arg in get_args(tp)) + return f"typing.Literal[{args}]" module = getattr(tp, "__module__", None) qualname = getattr(tp, "__qualname__", None) or getattr(tp, "__name__", None) if module and qualname: @@ -772,6 +824,56 @@ def from_model(cls, model_cls: Type[Any]) -> SocketSpec: "from_model expects a Pydantic BaseModel subclass, dataclass, or TypedDict." ) + @classmethod + def _leaf_from_literal(cls, T: Any) -> SocketSpec: + """Build the leaf spec for ``Literal[...]``. + + The socket keeps the permitted values in the ``allowed_values`` extra. + ``structured_type`` is added when one ``Enum`` supplies every argument, + whether the argument was written as a member or as that member's value, + so the member is rebuilt after serialization; arguments spanning two + enums, or naming a value no member of the one enum carries, keep their + values alone. ``literal_base`` records the type such a socket widens + into; an enum socket records none, because its output side carries the + member rather than the bare value. An argument no spec round trip can + carry (``bytes``) leaves the socket unconstrained, as an unrecognized + annotation always has. + """ + args = list(get_args(T)) + py_type = cls._py_type_name(T) + if not args: + return SocketSpec(identifier=cls.DEFAULT, meta=SocketMeta()) + + enum_classes = {type(arg) for arg in args if isinstance(arg, Enum)} + values: list = [] + for arg in args: + value = literal_value(arg) + if not value_is_allowed(value, values): + values.append(value) + if not all(isinstance(value, JSON_SAFE_LITERALS) for value in values): + return SocketSpec( + identifier=cls.ANNOTATED, meta=SocketMeta(extras={"py_type": py_type}) + ) + + extras: Dict[str, Any] = {"py_type": py_type, "allowed_values": values} + enum_cls = next(iter(enum_classes)) if len(enum_classes) == 1 else None + member_values = [member.value for member in enum_cls] if enum_cls else [] + if enum_cls is not None and all( + value_is_allowed(value, member_values) for value in values + ): + extras["structured_type"] = structured_type_info(enum_cls) + else: + base = cls._literal_base_type(values) + if base is not None: + extras["literal_base"] = cls._leaf_from_type(base).identifier + return SocketSpec(identifier=cls.ANNOTATED, meta=SocketMeta(extras=extras)) + + @staticmethod + def _literal_base_type(values: list) -> Any: + """Return the single type behind ``values``, or ``None`` if they differ.""" + types_seen = {type(value) for value in values} + return next(iter(types_seen)) if len(types_seen) == 1 else None + @classmethod def _leaf_from_type(cls, T: Any) -> SocketSpec: if T is Any or T is inspect._empty or T is object: @@ -780,7 +882,20 @@ def _leaf_from_type(cls, T: Any) -> SocketSpec: if _is_typeddict_type(T): return cls._leaf_from_type(dict) + if is_enum_type(T): + return SocketSpec( + identifier=cls.ANNOTATED, + meta=SocketMeta( + extras={ + "py_type": cls._py_type_name(T), + "structured_type": structured_type_info(T), + } + ), + ) + origin = get_origin(T) + if origin is Literal: + return cls._leaf_from_literal(T) if _is_union_origin(origin): args = [] for arg in get_args(T): @@ -1050,10 +1165,16 @@ def build_inputs_from_signature( ): info = structured_type_info(base_T) if info is not None and "structured_type" not in spec.meta.extras: + # ``required=None`` so this overlay says nothing about + # requiredness: ``SocketMeta.required`` defaults to ``True`` + # and ``merge_meta`` prefers any non-``None`` overlay value, + # so a bare ``SocketMeta`` would republish ``required=True`` + # over the value computed from the parameter's default. spec = replace( spec, meta=merge_meta( - spec.meta, SocketMeta(extras={"structured_type": info}) + spec.meta, + SocketMeta(required=None, extras={"structured_type": info}), ), ) diff --git a/src/node_graph/task.py b/src/node_graph/task.py index e8b60a3..52a5bc9 100644 --- a/src/node_graph/task.py +++ b/src/node_graph/task.py @@ -410,8 +410,18 @@ def execute(self): """Execute the task.""" executor = self.resolve_executor_callable(require=True) inputs = self.inputs._value - args = [inputs[arg] for arg in self.args_data["args"]] - kwargs = {key: inputs[key] for key in self.args_data["kwargs"]} + positional = list(self.args_data["args"]) + keyword = list(self.args_data["kwargs"]) + if all(name in inputs for name in positional): + args = [inputs[name] for name in positional] + else: + # A socket nothing was written into collects no value, and leaving + # one out of a positional call would shift the rest; everything the + # task does hold goes by name, and what is missing is the + # callable's own to answer for. + args = [] + keyword = positional + keyword + kwargs = {key: inputs[key] for key in keyword if key in inputs} var_kwargs = ( inputs[self.args_data["var_kwargs"]] if self.args_data["var_kwargs"] @@ -463,7 +473,10 @@ def set_inputs( or "property" to force property values even when linked. """ + from node_graph.input_model import validate_task_inputs + data = deep_copy_only_dicts(data) + validate_task_inputs(self, data) if value_source is None: value_source = "property" if self._allow_input_overrides else "link" if value_source not in ("link", "property"): diff --git a/src/node_graph/task_spec.py b/src/node_graph/task_spec.py index 349de46..9243afe 100644 --- a/src/node_graph/task_spec.py +++ b/src/node_graph/task_spec.py @@ -239,6 +239,23 @@ def _prepare_call_inputs(self, exec_obj, *args, **kwargs): ) return dict(kwargs) + def _validate_call_inputs(self, exec_obj, inputs) -> None: + """Check what is written at this call against the task's input model. + + A socket passes: it stands for a value nobody has yet. Every value + beside it is checked against the field it is written to -- a tagged + value by what it carries -- so a typo in a graph body fails at the + line that wrote it rather than after the graph has been submitted. + """ + from node_graph.input_model import ( + input_model_of_callable, + validate_wiring_inputs, + ) + + model = input_model_of_callable(exec_obj) + if model is not None: + validate_wiring_inputs(model, inputs, label=self.identifier.split(".")[-1]) + def __call__(self, *args, **kwargs): graph = self._get_current_graph() @@ -255,6 +272,7 @@ def __call__(self, *args, **kwargs): if isinstance(exec_obj, BaseHandle) and hasattr(exec_obj, "_callable"): exec_obj = exec_obj._callable prepared_inputs = self._prepare_call_inputs(exec_obj, *args, **kwargs) + self._validate_call_inputs(exec_obj, prepared_inputs) task.set_inputs(prepared_inputs) diff --git a/src/node_graph/tasks/function_task.py b/src/node_graph/tasks/function_task.py index b923890..d88ea45 100644 --- a/src/node_graph/tasks/function_task.py +++ b/src/node_graph/tasks/function_task.py @@ -32,6 +32,7 @@ def build( error_handlers: Optional[Dict[str, ErrorHandlerSpec]] = None, metadata: Optional[dict] = None, version: Optional[str] = None, + executor: Optional[Callable] = None, ) -> TaskSpec: """ - infers function I/O @@ -46,7 +47,10 @@ def build( func_in, func_out = infer_specs_from_callable(obj, input_spec, output_spec) error_handlers = normalize_error_handlers(error_handlers) metadata = dict(metadata or {}) - executor = RuntimeExecutor.from_callable(obj) + # The spec is inferred from ``obj``, so its signature, source and + # return annotation stay visible even when something else runs. + run_callable = executor or obj + executor = RuntimeExecutor.from_callable(run_callable) # We always use the EMBEDDED schema for the function task, but when storing the spec in the DB, # we will check if the callable is a BaseHandler, and switch the schema_source to HANDLER accordingly. # This avoids cyclic import. @@ -66,5 +70,5 @@ def build( ) is_graph_task = str(task_type).upper() == "GRAPH" handle = GraphTaskHandle(spec) if is_graph_task else TaskHandle(spec) - handle._callable = obj + handle._callable = run_callable return handle diff --git a/src/node_graph/utils/__init__.py b/src/node_graph/utils/__init__.py index d0507f7..ae98f38 100644 --- a/src/node_graph/utils/__init__.py +++ b/src/node_graph/utils/__init__.py @@ -105,6 +105,24 @@ def resolve_tagged_values(inputs: Any) -> Any: return inputs +def untagged_copy(inputs: Any) -> Any: + """Return a copy of ``inputs`` with every ``TaggedValue`` unwrapped. + + ``resolve_tagged_values`` unwraps in place, which is right when the caller + is done with the tags and wrong when it is not: a graph body still needs + them to draw its links. + """ + from node_graph.socket import TaggedValue + + if isinstance(inputs, TaggedValue): + return untagged_copy(inputs.__wrapped__) + if isinstance(inputs, dict): + return {key: untagged_copy(value) for key, value in inputs.items()} + if isinstance(inputs, (list, tuple)): + return type(inputs)(untagged_copy(value) for value in inputs) + return inputs + + def get_arg_type(name: str, args_data: dict, arg_type: str = "kwargs") -> None: """Get the argument type from the input data.""" if arg_type.upper() == "ARGS": diff --git a/src/node_graph/utils/graph.py b/src/node_graph/utils/graph.py index ca344e9..eb53ffd 100644 --- a/src/node_graph/utils/graph.py +++ b/src/node_graph/utils/graph.py @@ -134,6 +134,60 @@ def _assign_graph_outputs(outputs: Any, graph: Graph) -> None: ) +def _deserialize_inputs(namespace: Any, values: Any, adapter: Any) -> Any: + """Recursively apply ``adapter.deserialize`` to leaves of ``values`` so a + ``@task.graph`` body receives the primitive its signature declares, even + when the stored value was wrapped by the engine for provenance.""" + from node_graph.socket import TaskSocketNamespace + + if adapter is None or not hasattr(adapter, "deserialize"): + return values + if not isinstance(values, dict): + return values + from node_graph.utils.struct_utils import is_structured_instance + + out = dict(values) + for name, item in namespace._sockets.items(): + if name not in out: + continue + value = out[name] + if isinstance(item, TaskSocketNamespace): + if isinstance(value, dict): + out[name] = _deserialize_inputs(item, value, adapter) + elif is_structured_instance(value): + # The namespace was already materialised into a dataclass / + # Pydantic instance by ``coerce_inputs_from_spec`` (which + # runs before ``_deserialize_inputs`` in ``materialize_graph``). + # Without this branch, the adapter never sees the structured + # value, so a serialiser that auto-promotes primitive fields + # to engine-typed wrappers (e.g. ``aiida-workgraph``'s + # ``orm.Int`` / ``orm.Float``) leaves them wrapped inside + # the dataclass and downstream ``int``-typed code breaks. + out[name] = adapter.deserialize(value, item) + else: + out[name] = adapter.deserialize(value, item) + return out + + +def _validate_graph_body_inputs(func: Callable, inputs: dict, name: str) -> dict: + """Return a graph's resolved inputs, checked against its input model. + + Every expansion of a ``@task.graph`` passes through here -- ``build()`` on + a handle, the engine building a subgraph, and the workgraph task running + one -- so this is the one place a graph's whole payload is known. + + A graph with no model gets its inputs back untouched; one with a model + gets them in the types the model declares, still under the tags its links + are drawn from. + """ + from node_graph.input_model import input_model_of_callable, validate_graph_inputs + + model = input_model_of_callable(func) + if model is None: + return inputs + return validate_graph_inputs(model, inputs, label=name) + + def materialize_graph( func: Callable, in_spec: SocketSpec, @@ -173,6 +227,10 @@ def materialize_graph( tag_socket_value(graph.inputs) inputs = graph.inputs._collect_values(unwrap=False) inputs = coerce_inputs_from_spec(inputs, in_spec) + inputs = _deserialize_inputs( + graph.inputs, inputs, getattr(graph, "serialization", None) + ) + inputs = _validate_graph_body_inputs(func, inputs, name) raw = func(**inputs) _assign_graph_outputs(raw, graph) tag_socket_value(graph.inputs, only_uuid=True) diff --git a/src/node_graph/utils/struct_utils.py b/src/node_graph/utils/struct_utils.py index 3381b17..9405b9b 100644 --- a/src/node_graph/utils/struct_utils.py +++ b/src/node_graph/utils/struct_utils.py @@ -1,11 +1,19 @@ from __future__ import annotations from dataclasses import asdict, is_dataclass +from enum import Enum from typing import Any, Dict import importlib from pydantic import BaseModel +from node_graph.errors import SocketValueError + + +def is_enum_type(tp: Any) -> bool: + """Return True for Enum subclasses (including ``str``- and ``int``-Enum).""" + return isinstance(tp, type) and issubclass(tp, Enum) + def is_structured_instance(value: Any) -> bool: """Return True for dataclass or Pydantic model instances.""" @@ -29,11 +37,13 @@ def structured_to_dict(value: Any) -> Any: def structured_type_info(tp: Any) -> Dict[str, str] | None: - """Return a serializable descriptor for structured types.""" + """Return a serializable descriptor for structured types (incl. ``Enum``).""" if is_dataclass(tp): return {"kind": "dataclass", "path": structured_type_path(tp)} if isinstance(tp, type) and issubclass(tp, BaseModel): return {"kind": "pydantic", "path": structured_type_path(tp)} + if is_enum_type(tp): + return {"kind": "enum", "path": structured_type_path(tp)} return None @@ -50,16 +60,233 @@ def import_structured_type(path: str) -> Any: return current -def coerce_structured_value(value: Any, info: Dict[str, str] | None) -> Any: - """Rebuild a structured instance from a dict when spec says so.""" +#: Literal arguments a spec can carry through a to_dict/from_dict round trip. +JSON_SAFE_LITERALS = (str, bool, int, type(None)) + + +def untagged(value: Any) -> Any: + """Return the value a ``TaggedValue`` wraps, else ``value``. + + Canonicalization builds a fresh object, so the proxy has to come off first: + ``isinstance`` reports the wrapped type and would otherwise pass a proxy on + where callers expect a plain value. + """ + from node_graph.socket import TaggedValue + + return value.__wrapped__ if isinstance(value, TaggedValue) else value + + +def retagged(value: Any, original: Any) -> Any: + """Return ``value`` carrying ``original``'s tag, when it had one. + + Rebuilding a value loses the ``TaggedValue`` a socket wrapped it in, and + with it the socket the graph body needs to raise a link instead of a + literal. The uuid is carried over so provenance still points at one value. + """ + from node_graph.socket import TaggedValue + + if not isinstance(original, TaggedValue): + return value + tagged = TaggedValue(value, socket=original._socket) + tagged._self_uuid = original._uuid + return tagged + + +def literal_value(value: Any) -> Any: + """Return the bare value an ``Enum`` member stands for, else ``value``.""" + return value.value if isinstance(value, Enum) else value + + +def value_is_allowed(value: Any, allowed: Any) -> bool: + """Return True when ``value`` is one of ``allowed``. + + Two numbers match only when their types agree, which is typing's rule that + ``True`` is not ``1`` and ``1`` is not ``1.0``. Anything else matches on + equality alone, so a value that arrives wrapped, such as a storage node + holding ``'none'``, still names the value it equals. + """ + return any(_values_match(item, value) for item in allowed) + + +def _values_match(candidate: Any, value: Any) -> bool: + """Return True when ``value`` is the same value as ``candidate``.""" + numbers = (bool, int, float, complex) + if isinstance(candidate, numbers) and isinstance(value, numbers): + return type(candidate) is type(value) and candidate == value + try: + return bool(candidate == value) + except (TypeError, ValueError): + # An equality that yields an array, or none at all, decides nothing. + return False + + +def format_allowed_values(allowed: Any) -> str: + """Render allowed values as ``'a', 'b' or 'c'``.""" + rendered = [repr(item) for item in allowed] + if not rendered: + return "" + if len(rendered) == 1: + return rendered[0] + return f"{', '.join(rendered[:-1])} or {rendered[-1]}" + + +def socket_subject(name: str) -> str: + """Render the phrase an error uses to point at the socket ``name``.""" + return f"socket '{name}'" + + +def socket_loc(full_name: str | None, *keys: str) -> tuple: + """Return the path to a socket, spelled as an input model spells its errors. + + ``full_name`` runs from the task's ``inputs`` or ``outputs``, which the + model's own path does not name, so the first segment is dropped. ``keys`` + are the names below the socket that the value itself carries. + """ + parts = [part for part in (full_name or "").split(".") if part] + if parts and parts[0] in ("inputs", "outputs"): + parts = parts[1:] + return tuple(parts) + keys + + +def _invalid_value_error( + value: Any, + allowed: Any, + subject: str | None, + hint: str, + *, + loc: tuple = (), + error_type: str = "literal_error", +) -> SocketValueError: + return SocketValueError( + f"Invalid value for {subject or 'this input'}.\n" + f" Input should be {format_allowed_values(allowed)}. Got {value!r}.\n" + f" {hint}", + loc=loc, + error_type=error_type, + ) + + +def canonical_enum_member( + value: Any, + cls: type, + *, + allowed: Any = None, + subject: str | None = None, + loc: tuple = (), +) -> Any: + """Return the ``cls`` member ``value`` names, raising if it names none. + + A member of ``cls``, a member of any other ``Enum`` carrying the same + value, and a bare member value all name the same member: membership is + decided by ``value_is_allowed`` over the members' values, because + serialization keeps only the value. That is the rule a ``Literal`` socket + applies too, so an ``IntEnum`` whose members are ``1`` and ``2`` rejects + ``True`` exactly as ``Literal[1, 2]`` does. ``allowed``, when given, + restricts the result to that subset of member values. + """ + permitted = [member.value for member in cls] if allowed is None else list(allowed) + value = untagged(value) + if isinstance(value, cls): + member = value + else: + candidate = literal_value(value) + member = next( + (item for item in cls if value_is_allowed(item.value, [candidate])), + None, + ) + if member is None: + example = next(iter(cls.__members__), None) + named = f" ({cls.__name__}.{example})" if example else "" + raise _invalid_value_error( + value, + permitted, + subject, + f"{structured_type_path(cls)} members are accepted by " + f"member{named} or by value.", + loc=loc, + error_type="enum", + ) + if not value_is_allowed(member.value, permitted): + raise _invalid_value_error( + value, + permitted, + subject, + f"The socket admits only part of {structured_type_path(cls)}.", + loc=loc, + error_type="literal_error", + ) + return member + + +def canonical_literal_value( + value: Any, + allowed: Any, + *, + subject: str | None = None, + loc: tuple = (), +) -> Any: + """Return ``value`` when it is one of ``allowed``, raising otherwise. + + Matching follows typing's rule that ``True`` is not ``1``: a value matches + only a candidate of its own type. An ``Enum`` member matches by its value. + """ + candidate = literal_value(untagged(value)) + if value_is_allowed(candidate, allowed): + return candidate + raise _invalid_value_error( + value, + allowed, + subject, + "Only the values listed above are accepted here.", + loc=loc, + error_type="literal_error", + ) + + +def canonical_socket_value( + value: Any, + *, + structured_type: Dict[str, str] | None = None, + allowed: Any = None, + subject: str | None = None, + loc: tuple = (), +) -> Any: + """Return the value a socket may store, raising when the socket forbids it. + + ``value`` comes back untouched when the socket constrains nothing. + """ + if structured_type is not None and structured_type.get("kind") == "enum": + cls = import_structured_type(structured_type["path"]) + return canonical_enum_member( + value, cls, allowed=allowed, subject=subject, loc=loc + ) + if allowed is not None: + return canonical_literal_value(value, allowed, subject=subject, loc=loc) + return value + + +def coerce_structured_value( + value: Any, info: Dict[str, str] | None, name: str | None = None +) -> Any: + """Rebuild a structured instance from a flat value when spec says so.""" if info is None: return value + kind = info.get("kind") + # Enum check goes before the dict gate: its serialized form is a bare value, + # so an already-materialized-instance/dict short-circuit would never fire. + # Assignment already decided membership, so this only rebuilds the member. + if kind == "enum": + cls = import_structured_type(info["path"]) + subject = socket_subject(name) if name else None + return retagged(canonical_enum_member(value, cls, subject=subject), value) if is_structured_instance(value): return value if not isinstance(value, dict): return value + # Only import the target class once we know we must rebuild from a dict. + # Imports of already-materialized instances (e.g. locally-defined types + # whose qualname is not importable) must not reach here. cls = import_structured_type(info["path"]) - kind = info.get("kind") if kind == "pydantic": if contains_tagged_value(value): if hasattr(cls, "model_construct"): @@ -88,7 +315,7 @@ def coerce_inputs_from_spec(values: Any, spec: Any) -> Any: continue info = child.meta.extras.get("structured_type") if info: - out[name] = coerce_structured_value(out[name], info) + out[name] = coerce_structured_value(out[name], info, name) continue if child.is_namespace() and isinstance(out[name], dict): out[name] = coerce_inputs_from_spec(out[name], child) diff --git a/tests/test_engine_local.py b/tests/test_engine_local.py index 76a3037..e3f06f7 100644 --- a/tests/test_engine_local.py +++ b/tests/test_engine_local.py @@ -7,6 +7,7 @@ from node_graph.engine.provenance import ProvenanceRecorder from typing import Annotated, Any from dataclasses import dataclass +from enum import Enum from pydantic import BaseModel @@ -213,3 +214,56 @@ def test_local_engine_accepts_dataclass_models(): assert results["sum"] == 7 assert results["product"] == 12 + + +class Color(str, Enum): + RED = "red" + GREEN = "green" + BLUE = "blue" + + +class Priority(int, Enum): + LOW = 1 + HIGH = 9 + + +@task(outputs=ns(name=str)) +def color_name(color: Color) -> dict: + assert isinstance(color, Color), f"expected Color, got {type(color).__name__}" + return {"name": color.name} + + +@task(outputs=ns(name=str)) +def priority_name(priority: Priority) -> dict: + assert isinstance( + priority, Priority + ), f"expected Priority, got {type(priority).__name__}" + return {"name": priority.name} + + +def test_local_engine_coerces_str_enum(): + # A ``str``-Enum member flattens to a bare ``str`` across a serialization + # boundary (it *is* a str). Feed the engine that flattened form, not the + # member itself, so the coercion has something real to rebuild; passing + # ``Color.GREEN`` directly would leave the body's isinstance check trivially + # satisfied and the test unable to fail. The signature records the enum type, + # so ``coerce_inputs_from_spec`` must turn ``"green"`` back into ``Color.GREEN`` + # before the body runs. + ng = Graph(name="local-enum", outputs=ns(name=str)) + node = ng.add_task(color_name, "pick", color="green") + ng.add_link(node.outputs.name, ng.outputs.name) + + results = LocalEngine().run(ng) + + assert results["name"] == "GREEN" + + +def test_local_engine_coerces_int_enum(): + # Same for an ``int``-Enum: its flattened form is a bare ``int``. + ng = Graph(name="local-int-enum", outputs=ns(name=str)) + node = ng.add_task(priority_name, "pick", priority=9) + ng.add_link(node.outputs.name, ng.outputs.name) + + results = LocalEngine().run(ng) + + assert results["name"] == "HIGH" diff --git a/tests/test_enum_literal_sockets.py b/tests/test_enum_literal_sockets.py new file mode 100644 index 0000000..8a59971 --- /dev/null +++ b/tests/test_enum_literal_sockets.py @@ -0,0 +1,464 @@ +"""An enum or Literal socket accepts what names one of its members, nothing else.""" + +from enum import Enum, IntEnum +from typing import Literal, Optional + +import pytest + +from node_graph import Graph, task +from node_graph.engine.local import LocalEngine +from node_graph.socket_spec import SocketSpecAPI as api + + +class Spin(Enum): + NONE = "none" + COLLINEAR = "collinear" + SPIN_ORBIT = "spin_orbit" + + +class Narrow(Enum): + """A separate enum whose members repeat two of ``Spin``'s values.""" + + NONE = "none" + COLLINEAR = "collinear" + + +class NameOnly(Enum): + """Shares a member NAME with ``Narrow``, but not its value.""" + + NONE = "something_else" + + +class Priority(IntEnum): + """An enum whose members carry bare ints.""" + + LOW = 1 + HIGH = 2 + + +@task() +def consume_narrow(spin: Narrow) -> str: + assert isinstance(spin, Narrow), f"expected Narrow, got {type(spin).__name__}" + return spin.value + + +@task.graph() +def narrow_graph(spin: Narrow) -> str: + return consume_narrow(spin=spin).result + + +@task.graph() +def defer_literal(spin) -> str: + """Call the enum-typed sub-graph with whatever the parent was given.""" + return narrow_graph(spin=spin) + + +ACCEPTED = [ + pytest.param(Narrow.NONE, id="declared-member"), + pytest.param(Spin.NONE, id="foreign-member-same-value"), + pytest.param("none", id="bare-value"), +] + +REJECTED = [ + pytest.param("banana", id="unknown-value"), + pytest.param(Spin.SPIN_ORBIT, id="foreign-member-unknown-value"), + pytest.param(NameOnly.NONE, id="foreign-member-matching-name-only"), + pytest.param(42, id="wrong-type"), +] + + +@pytest.mark.parametrize("value", ACCEPTED) +def test_entry_graph_accepts_anything_naming_a_member(value): + ng = narrow_graph.build(spin=value) + + # Storage keeps the member's value, the form that survives a process + # boundary; the task body is handed the member back. + assert ng.inputs.spin.value == Narrow.NONE.value + + +@pytest.mark.parametrize("value", REJECTED) +def test_entry_graph_rejects_a_value_naming_no_member(value): + with pytest.raises(ValueError, match="Input should be 'none' or 'collinear'"): + narrow_graph.build(spin=value) + + +@pytest.mark.parametrize("value", ACCEPTED) +def test_deferred_subgraph_accepts_anything_naming_a_member(value): + @task.graph() + def outer() -> str: + return narrow_graph(spin=value) + + ng = outer.build() + + assert ng.tasks.narrow_graph.inputs.spin.value == Narrow.NONE.value + + +@pytest.mark.parametrize("value", REJECTED) +def test_deferred_subgraph_rejects_at_build_not_at_run(value): + @task.graph() + def outer() -> str: + return narrow_graph(spin=value) + + with pytest.raises(ValueError, match="Input should be 'none' or 'collinear'"): + outer.build() + + +@pytest.mark.parametrize("value", REJECTED) +def test_link_from_untyped_parent_socket_rejects_at_build(value): + """The parent socket is untyped, so the link is where the value is read.""" + with pytest.raises(ValueError, match="Input should be 'none' or 'collinear'"): + defer_literal.build(spin=value) + + +def test_error_names_the_socket_and_omits_the_value_wrapper(): + with pytest.raises(ValueError) as excinfo: + narrow_graph.build(spin="banana") + + message = str(excinfo.value) + assert "graph_inputs.inputs.spin" in message + assert "TaggedValue" not in message + assert "uuid" not in message + + +def test_body_receives_a_member_rebuilt_from_the_stored_value(): + """A run stores the bare value; the body still sees the member.""" + ng = Graph(name="enum-readback") + node = ng.add_task(consume_narrow, "pick", spin=Spin.COLLINEAR) + ng.outputs.result = node.outputs.result + + results = LocalEngine().run(ng) + + assert results["result"] == "collinear" + + +@pytest.mark.parametrize("value", ACCEPTED) +def test_an_enum_graph_input_reaches_the_task_as_a_link(value): + """The rebuilt member keeps the socket, so the body wires rather than copies.""" + ng = narrow_graph.build(spin=value) + + sources = {link.from_task.name for link in ng.links} + + assert "graph_inputs" in sources + + +def test_an_enum_parameter_keeps_the_requiredness_of_its_default(): + def signature(a: Spin, b: Spin = Spin.NONE, c: Optional[Spin] = None): + ... + + fields = api.build_inputs_from_signature(signature).fields + + assert fields["a"].meta.required is True + assert fields["b"].meta.required is False + assert fields["b"].default == Spin.NONE.value + assert fields["c"].meta.required is False + + +# --- defaults -------------------------------------------------------------- + + +@task() +def consume_defaulted(spin: Narrow = Narrow.COLLINEAR) -> str: + assert isinstance(spin, Narrow), f"expected Narrow, got {type(spin).__name__}" + return spin.value + + +def test_a_default_is_stored_as_the_value_an_assignment_would_store(): + ng = Graph(name="defaults") + defaulted = ng.add_task(consume_defaulted, "defaulted") + assigned = ng.add_task(consume_defaulted, "assigned", spin=Narrow.COLLINEAR) + + assert defaulted.inputs.spin.value == assigned.inputs.spin.value + assert defaulted.inputs.spin._serialize_value() == Narrow.COLLINEAR.value + + +def test_a_defaulted_socket_reads_the_same_after_a_graph_round_trip(): + ng = Graph(name="defaults-round-trip") + ng.add_task(consume_defaulted, "defaulted") + + restored = Graph.from_dict(ng.to_dict()) + + assert restored.tasks.defaulted.inputs.spin.value == Narrow.COLLINEAR.value + + +def test_a_defaulted_socket_still_hands_the_body_a_member(): + ng = Graph(name="defaults-body") + node = ng.add_task(consume_defaulted, "defaulted") + ng.outputs.result = node.outputs.result + + assert LocalEngine().run(ng)["result"] == "collinear" + + +def test_a_default_naming_no_member_is_rejected_where_the_spec_is_built(): + with pytest.raises(ValueError, match="Input should be 'none' or 'collinear'"): + + @task() + def bad_default(spin: Narrow = Spin.SPIN_ORBIT) -> str: + return "unreachable" + + +def test_a_literal_default_outside_the_allowed_set_is_rejected(): + """The value a build rejects is rejected as a default too.""" + with pytest.raises(ValueError, match="Input should be 'none' or 'collinear'"): + + @task.graph() + def bad_literal_default(spin: Literal["none", "collinear"] = "banana") -> str: + return consume_narrow(spin=Narrow.NONE).result + + +def test_optional_enum_socket_still_takes_none(): + @task.graph() + def maybe(spin: Narrow = None) -> str: + return consume_narrow(spin=Narrow.NONE).result + + ng = maybe.build(spin=None) + + assert ng.inputs.spin.value is None + + +# --- Literal --------------------------------------------------------------- + + +def test_literal_of_enum_members_carries_the_enum_and_the_subset(): + spec = api._leaf_from_type(Literal[Spin.NONE, Spin.COLLINEAR]) + + assert spec.meta.extras["allowed_values"] == ["none", "collinear"] + assert spec.meta.extras["structured_type"]["kind"] == "enum" + assert spec.meta.extras["structured_type"]["path"].endswith(".Spin") + + +def test_literal_of_strings_carries_its_base_type(): + spec = api._leaf_from_type(Literal["none", "collinear"]) + + assert spec.meta.extras["allowed_values"] == ["none", "collinear"] + assert spec.meta.extras["literal_base"] == api._leaf_from_type(str).identifier + assert "structured_type" not in spec.meta.extras + + +def test_literal_of_mixed_types_keeps_the_values_without_a_base_type(): + spec = api._leaf_from_type(Literal["a", 1]) + + assert spec.meta.extras["allowed_values"] == ["a", 1] + assert "literal_base" not in spec.meta.extras + + +def test_literal_of_unrepresentable_arguments_constrains_nothing(): + spec = api._leaf_from_type(Literal[b"x"]) + + assert "allowed_values" not in spec.meta.extras + + +def test_a_literal_socket_survives_a_spec_round_trip(): + from node_graph.socket_spec import SocketSpec + + spec = api._leaf_from_type(Literal[Spin.NONE, Spin.COLLINEAR]) + + restored = SocketSpec.from_dict(spec.to_dict()) + + assert restored.meta.extras == spec.meta.extras + + +def test_literal_naming_one_enum_by_value_still_carries_the_enum(): + """``Spin.NONE`` and ``'collinear'`` name members of the same enum.""" + spec = api._leaf_from_type(Literal[Spin.NONE, "collinear"]) + + assert spec.meta.extras["allowed_values"] == ["none", "collinear"] + assert spec.meta.extras["structured_type"]["path"].endswith(".Spin") + + +def test_literal_naming_a_value_no_member_carries_keeps_the_values_alone(): + spec = api._leaf_from_type(Literal[Spin.NONE, "banana"]) + + assert spec.meta.extras["allowed_values"] == ["none", "banana"] + assert "structured_type" not in spec.meta.extras + + +def test_py_type_name_distinguishes_two_literals(): + assert api._py_type_name(Literal[1, 2]) != api._py_type_name(Literal["a", "b"]) + assert api._py_type_name(Literal["a", "b"]) == "typing.Literal['a', 'b']" + + +def test_py_type_name_distinguishes_same_named_enums_from_two_modules(): + here = Enum("Duplicate", {"NONE": "none"}, module="package_a") + there = Enum("Duplicate", {"NONE": "none"}, module="package_b") + + assert api._py_type_name(Literal[here.NONE]) != api._py_type_name( + Literal[there.NONE] + ) + + +def test_an_intenum_socket_and_a_literal_of_its_values_agree_on_true(): + """``True`` is not ``1`` on either path.""" + from node_graph.utils.struct_utils import canonical_socket_value + + info = api._leaf_from_type(Priority).meta.extras["structured_type"] + assert canonical_socket_value(1, structured_type=info) is Priority.LOW + for rejected in (True, 1.0): + with pytest.raises(ValueError, match="Input should be 1 or 2"): + canonical_socket_value(rejected, structured_type=info) + with pytest.raises(ValueError, match="Input should be 1 or 2"): + canonical_socket_value(rejected, allowed=[1, 2]) + + +def test_literal_socket_takes_a_member_of_its_subset_by_value(): + @task.graph() + def only_none(spin: Literal[Spin.NONE]) -> str: + return consume_narrow(spin=Narrow.NONE).result + + assert only_none.build(spin="none").inputs.spin.value == Spin.NONE.value + assert only_none.build(spin=Narrow.NONE).inputs.spin.value == Spin.NONE.value + with pytest.raises(ValueError, match="Input should be 'none'"): + only_none.build(spin=Spin.COLLINEAR) + + +# --- links ----------------------------------------------------------------- + + +@task() +def emit_subset() -> Literal[Spin.NONE, Spin.COLLINEAR]: + return Spin.NONE + + +@task() +def emit_spin() -> Spin: + return Spin.NONE + + +@task() +def emit_ints() -> Literal[1, 2]: + return 1 + + +@task() +def emit_one_string() -> Literal["none"]: + return "none" + + +@task() +def take_spin(x: Spin) -> int: + return 1 + + +@task() +def take_subset(x: Literal[Spin.NONE, Spin.COLLINEAR]) -> int: + return 1 + + +@task() +def take_strings(x: Literal["none", "collinear"]) -> int: + return 1 + + +@task() +def take_str(x: str) -> int: + return 1 + + +def link(producer, consumer): + ng = Graph(name="link") + src = ng.add_task(producer, "src") + dst = ng.add_task(consumer, "dst") + ng.add_link(src.outputs.result, dst.inputs.x) + + +def test_a_subset_flows_into_the_whole_enum(): + link(emit_subset, take_spin) + + +def test_the_whole_enum_does_not_flow_into_a_subset(): + with pytest.raises(TypeError, match="Socket value range mismatch"): + link(emit_spin, take_subset) + + +def test_a_subset_flows_into_another_enum_carrying_the_same_values(): + """Membership is by value, so the declaring class need not match.""" + + @task() + def take_narrow_enum(x: Narrow) -> int: + return 1 + + link(emit_subset, take_narrow_enum) + + +def test_literals_of_different_base_types_do_not_mix(): + with pytest.raises(TypeError, match="Socket value range mismatch"): + link(emit_ints, take_strings) + + +def test_a_string_subset_flows_into_a_wider_string_literal(): + link(emit_one_string, take_strings) + + +def test_a_string_literal_widens_into_its_base_type(): + link(emit_one_string, take_str) + + +def test_a_literal_of_enum_members_does_not_widen_into_the_value_type(): + """An enum output carries the member, which a ``str`` socket cannot take.""" + with pytest.raises(TypeError, match="Socket type mismatch"): + link(emit_subset, take_str) + + +def test_a_typed_source_of_the_same_base_does_not_flow_into_a_literal(): + """``str`` is not restricted to the two strings the target admits.""" + with pytest.raises(TypeError, match="Socket value range mismatch"): + link(take_str, take_strings) + + +def test_an_untyped_source_flows_into_a_literal_and_is_read_at_run(): + """An untyped hop is how a parent passes a value on, so the link stands.""" + + @task() + def emit_untyped(): + return "banana" + + link(emit_untyped, take_strings) + + +def test_a_value_two_untyped_hops_out_is_rejected_when_the_hop_expands(): + """The restricted socket appears only as the inner graph expands, at run.""" + + @task.graph() + def one_hop(spin) -> str: + return narrow_graph(spin=spin) + + @task.graph() + def two_hops(spin) -> str: + return one_hop(spin=spin) + + with pytest.raises(ValueError, match="Input should be 'none' or 'collinear'"): + LocalEngine().run(two_hops.build(spin=Spin.SPIN_ORBIT)) + + assert LocalEngine().run(two_hops.build(spin=Narrow.NONE))["result"] == "none" + + +def test_the_run_time_message_names_the_socket_that_refused_the_value(): + """An untyped output is read at run, where the name is the socket's own.""" + + @task() + def emit_foreign(): + return Spin.SPIN_ORBIT + + @task.graph() + def from_task() -> str: + return narrow_graph(spin=emit_foreign().result) + + with pytest.raises(ValueError, match="socket 'spin'"): + LocalEngine().run(from_task.build()) + + +def test_a_wrapped_value_read_back_from_storage_still_names_its_member(): + """Storage hands back a node that equals the value but is not a ``str``.""" + from node_graph.utils.struct_utils import canonical_socket_value + + class StoredString: + def __init__(self, value): + self.value = value + + def __eq__(self, other): + return self.value == other + + info = api._leaf_from_type(Narrow).meta.extras["structured_type"] + + assert canonical_socket_value(StoredString("none"), structured_type=info) is ( + Narrow.NONE + ) diff --git a/tests/test_input_model.py b/tests/test_input_model.py new file mode 100644 index 0000000..578dfea --- /dev/null +++ b/tests/test_input_model.py @@ -0,0 +1,3101 @@ +"""A Pydantic model as a task's wire contract. + +``input_model=`` makes a model declare a task's input sockets and hold the +inputs to its rules at three moments: the call that wires the task, the +expansion of a ``@task.graph``, and the run edge of a leaf task. Each test +below is one claim about that contract. +""" + +from __future__ import annotations + +import enum +from decimal import Decimal +from typing import Annotated, Any, Literal, Optional, Union + +import pytest +from pydantic import ( + AfterValidator, + BaseModel, + ConfigDict, + Field, + PlainSerializer, + field_serializer, + ValidationError, + field_validator, + model_validator, +) +from typing_extensions import TypedDict + +from node_graph import Graph, task +from node_graph.engine.local import LocalEngine +from node_graph.input_model import ( + BODY_RECEIVES, + ModelContractError, + ModelDerivedValueError, + TaskInputValidationError, + TaskOutputValidationError, + dump_model_field, + spec_from_model, + validate_wiring_inputs, +) +from node_graph.socket_spec import namespace as ns + + +class Color(enum.Enum): + RED = "red" + BLUE = "blue" + + +def links_of(graph: Graph) -> list[str]: + """Return every link of ``graph`` as ``source -> target``, sorted.""" + return sorted( + f"{link.from_socket._full_name_with_task} -> {link.to_socket._full_name_with_task}" + for link in graph.links + ) + + +@pytest.fixture +def without_wiring_checks(monkeypatch): + """Run the body with checkpoint A disabled, as a negative control. + + Both of its seams are stubbed: the call, and the write into the task's + sockets that every route to an input goes through. + """ + from node_graph import input_model + from node_graph.task_spec import BaseHandle + + monkeypatch.setattr( + BaseHandle, "_validate_call_inputs", lambda self, exec_obj, inputs: None + ) + monkeypatch.setattr(input_model, "validate_task_inputs", lambda task, inputs: None) + + +# -------------------------------------------------------------------------- +# 1. The model is the socket contract +# -------------------------------------------------------------------------- + + +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_and_their_defaults_come_from_the_model(): + """The bare signature contributes nothing: types, defaults and requiredness are the model's.""" + fields = add._spec.inputs.fields + assert fields["x"].identifier == "node_graph.int" + assert fields["x"].meta.required is True + assert fields["y"].identifier == "node_graph.int" + assert fields["y"].meta.required is False + assert fields["y"].default == 7 + + +def test_an_omitted_input_runs_on_the_models_default(): + assert add.run(x=2) == 9 + + +class NudgeInputs(BaseModel): + """``by`` defaults to ``None``, which reads as a missing input unless the model is asked.""" + + x: int + by: Optional[int] = 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(): + """``from_model`` leaves every field required; the model says which ones are.""" + assert nudge._spec.inputs.fields["by"].meta.required is False + + +def test_only_a_typed_mapping_is_dynamic(): + """A bare ``dict`` and a ``list`` name no members, so they stay leaf sockets.""" + + class Shapes(BaseModel): + bare: dict = {} + ordered: list[int] = [] + anything: dict[str, Any] = {} + + fields = spec_from_model(Shapes).fields + assert fields["bare"].identifier == "node_graph.any" + assert fields["ordered"].identifier == "node_graph.any" + assert fields["ordered"].meta.dynamic is False + assert fields["anything"].meta.dynamic is True + + +def test_a_typed_mapping_field_names_its_members_type(): + """The keys are unknown at definition time; every member's sockets are not.""" + + class Block(BaseModel): + width: int + label: str + + class Blocks(BaseModel): + blocks: dict[str, Block] + + blocks = spec_from_model(Blocks).fields["blocks"] + assert blocks.meta.dynamic is True + assert set(blocks.item.fields) == {"width", "label"} + assert blocks.item.fields["width"].identifier == "node_graph.int" + + +def test_an_open_topped_model_is_refused(): + class OpenTopped(BaseModel): + model_config = ConfigDict(extra="allow") + x: int + + with pytest.raises(ModelContractError, match="declares no contract"): + spec_from_model(OpenTopped) + + +def test_a_nested_open_topped_model_is_refused_too(): + """A field the nested model admits has no socket, so it never reaches storage.""" + + class OpenInner(BaseModel): + model_config = ConfigDict(extra="allow") + x: int + + class HoldsIt(BaseModel): + inner: OpenInner + + with pytest.raises(ModelContractError, match="OpenInner sets extra='allow'"): + spec_from_model(HoldsIt) + + +def test_an_open_topped_model_inside_a_container_is_refused_too(): + """The walk follows containers, so hiding the model in a list changes nothing.""" + + class OpenInner(BaseModel): + model_config = ConfigDict(extra="allow") + x: int + + class HoldsThem(BaseModel): + rows: list[OpenInner] + + with pytest.raises(ModelContractError, match="OpenInner sets extra='allow'"): + spec_from_model(HoldsThem) + + +def test_a_mapping_keyed_by_anything_but_str_is_refused(): + class IntKeyed(BaseModel): + counts: dict[int, str] + + with pytest.raises(ModelContractError, match="keys must be str"): + spec_from_model(IntKeyed) + + +def test_no_class_path_is_imported_to_rebuild_a_value(): + """A moved class cannot break stored data: nothing reads a path out of the spec to load it. + + The class's name is still recorded, under ``py_type``, and is only ever + compared as a string -- by link type-checking and in error messages. The + ``structured_type`` descriptor, which names a class to import, is the one + a model contract has no use for: the model rebuilds the value instead. + """ + + class PaintInputs(BaseModel): + color: Color + + @task(input_model=PaintInputs) + def paint(color): + return color + + extras = paint._spec.inputs.fields["color"].meta.extras + assert "structured_type" not in extras + assert extras["py_type"].endswith("Color") + + +# -------------------------------------------------------------------------- +# 2. Model and signature must agree, loudly, at decoration time +# -------------------------------------------------------------------------- + + +def test_a_signature_default_is_refused(): + with pytest.raises(ModelContractError, match="defaults live in the model"): + + @task(input_model=AddInputs) + def bad(x, y=7): + return x + y + + +def test_a_field_no_parameter_names_is_refused(): + with pytest.raises( + ModelContractError, match="declares 'y', which 'bad' does not take" + ): + + @task(input_model=AddInputs) + def bad(x): + return x + + +def test_a_parameter_no_field_declares_is_refused(): + with pytest.raises( + ModelContractError, match="takes 'z', which AddInputs does not declare" + ): + + @task(input_model=AddInputs) + def bad(x, y, z): + return x + + +def test_an_annotation_that_contradicts_the_field_is_refused(): + with pytest.raises(ModelContractError, match="but AddInputs declares"): + + @task(input_model=AddInputs) + def bad(x: str, y): + return x + + +def test_var_kwargs_are_refused(): + with pytest.raises(ModelContractError, match="which no model field can name"): + + @task(input_model=AddInputs) + def bad(x, y, **rest): + return x + + +def test_declaring_the_inputs_twice_is_refused(): + with pytest.raises(ModelContractError, match="keep input_model and drop inputs"): + + @task(input_model=AddInputs, inputs=spec_from_model(AddInputs)) + def bad(x, y): + return x + + +def test_an_output_model_on_a_graph_is_refused(): + with pytest.raises(ModelContractError, match="not supported on @task.graph"): + + class Out(BaseModel): + result: int + + @task.graph(input_model=AddInputs, output_model=Out) # type: ignore[call-arg] + def bad(x, y): + return add(x=x, y=y) + + +# -------------------------------------------------------------------------- +# 3. Checkpoint A -- what is written at the call +# -------------------------------------------------------------------------- + + +class PriceInputs(BaseModel): + """A ``Decimal`` and a fixed-length tuple: two shapes a socket reads as ``any``.""" + + amount: Decimal + pair: tuple[int, int] + ratio: float = Field(gt=0) + + +@task(input_model=PriceInputs) +def price(amount, pair, ratio): + return amount + + +GOOD_PRICE = {"amount": Decimal("1.5"), "pair": (1, 2), "ratio": 0.5} + + +@pytest.mark.parametrize( + "field, value", + [ + ("amount", "not a number"), + ("pair", [1, 2, 3]), + ("ratio", -1.0), + ], + ids=[ + "decimal-from-nonsense", + "tuple-of-the-wrong-length", + "constraint-the-socket-cannot-see", + ], +) +def test_a_bad_literal_fails_at_the_line_that_wrote_it(field, value): + """A field whose socket identifier is ``any`` has no other build-time check.""" + payload = dict(GOOD_PRICE, **{field: value}) + + @task.graph() + def wires_a_bad_value(): + return price(**payload) + + with pytest.raises(TaskInputValidationError) as excinfo: + wires_a_bad_value.build() + assert "Task 'price'" in str(excinfo.value) + assert field in str(excinfo.value) + + +@pytest.mark.parametrize( + "field, value", + [ + ("amount", "not a number"), + ("pair", [1, 2, 3]), + ("ratio", -1.0), + ], + ids=[ + "decimal-from-nonsense", + "tuple-of-the-wrong-length", + "constraint-the-socket-cannot-see", + ], +) +def test_a_bad_literal_at_the_call_passes_without_checkpoint_a( + without_wiring_checks, field, value +): + """The negative control: the socket layer types these as ``any`` and lets them through.""" + payload = dict(GOOD_PRICE, **{field: value}) + + @task.graph() + def wires_a_bad_value(): + return price(**payload) + + assert wires_a_bad_value.build() is not None + + +def test_a_literal_the_socket_layer_also_refuses_gets_the_models_message(): + """Where both can see the value, the model's error is the one that arrives.""" + + @task.graph() + def wires_a_string(m): + return add(x="sixty", y=m) + + with pytest.raises(TaskInputValidationError) as excinfo: + wires_a_string.build(m=1) + assert "Task 'add'" in str(excinfo.value) + assert "\nx\n" in str(excinfo.value) + + +class BoxInputs(BaseModel): + cfg: dict[str, int] + n: int + + +@task(input_model=BoxInputs) +def box(cfg, n): + return n + + +def test_a_bad_literal_inside_a_container_names_its_path(): + @task.graph() + def wires_a_bad_member(m): + return box(cfg={"x": "nope"}, n=m) + + with pytest.raises(TaskInputValidationError, match=r"cfg\.x"): + wires_a_bad_member.build(m=1) + + +def test_a_reference_inside_a_container_is_not_a_bad_value(): + """A link written into one member of a mapping is what a graph is for.""" + + @task.graph() + def wires_a_reference(m): + return box(cfg={"x": m}, n=m) + + assert len(wires_a_reference.build(m=1).links) == 3 + + +@task() +def untyped(a, b): + return a + + +def test_wiring_a_reference_leaves_every_link_where_it_was(monkeypatch): + """Checkpoint A validates and discards, so the graph it builds is the graph without it. + + Pydantic strips the proxy a tagged value wears for nearly every field + type, so a validated copy would turn a link into a literal. This compares + the same graph built with the check and without it. + """ + + @task.graph() + def chain(m, n): + first = untyped(a=m, b=n) + return add(x=m, y=first.result) + + from node_graph.task_spec import BaseHandle + + with monkeypatch.context() as patched: + patched.setattr( + BaseHandle, "_validate_call_inputs", lambda self, exec_obj, inputs: None + ) + without = links_of(chain.build(m=1, n=2)) + assert links_of(chain.build(m=1, n=2)) == without + assert any("graph_inputs.inputs.m -> add.inputs.x" == link for link in without) + + +def test_validating_at_the_call_hands_back_the_very_objects_it_was_given(): + """The values are passed on by identity, which is what keeps a reference a reference.""" + + class Sentinel: + pass + + given = {"cfg": {"x": 1}, "n": 2} + before = {"cfg": id(given["cfg"]), "n": id(given["n"])} + validate_wiring_inputs(BoxInputs, given, label="box") + assert {name: id(value) for name, value in given.items()} == before + + +class Spin(enum.Enum): + """The four spin treatments a workflow accepts.""" + + NONE = "none" + COLLINEAR = "collinear" + NON_COLLINEAR = "non_collinear" + SPIN_ORBIT = "spin_orbit" + + +class PhInputs(BaseModel): + 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 + + +#: What each run of ``ph``'s body was handed, so a test can see it did not run. +ph_ran_with: list = [] + + +@task(input_model=PhInputs) +def ph(spin, structure): + ph_ran_with.append(spin) + return 11.7 + + +class EpsInputs(BaseModel): + spin: Spin = Spin.NONE + structure: str + + +@task.graph(input_model=EpsInputs) +def eps(spin, structure): + return ph(spin=spin, structure=structure).result + + +def test_a_rule_on_the_inner_task_refuses_the_wiring_that_breaks_it(): + """The graph takes every spin; the rule is ph's own and fires where ph is wired. + + The value is refused at the line inside ``eps`` that hands it to ``ph``, so + the graph is never built and no ``ph`` runs. + """ + ph_ran_with.clear() + with pytest.raises(TaskInputValidationError, match="noncollinear magnetism"): + eps.build(spin=Spin.NON_COLLINEAR, structure="si") + assert ph_ran_with == [] + + +def test_a_spin_the_rule_admits_still_reaches_the_body_as_the_enum(): + """The control: the same wiring builds, runs, and hands the body the enum.""" + ph_ran_with.clear() + graph = eps.build(spin=Spin.COLLINEAR, structure="si") + assert "ph" in graph.tasks + graph.run() + assert ph_ran_with == [Spin.COLLINEAR] + + +@task() +def a_thousand(): + return 1000 + + +class Capped(BaseModel): + a: int + b: int + + @field_validator("a") + @classmethod + def cap(cls, value): + if value > 100: + raise ValueError("a must be at most 100") + return value + + +@task(input_model=Capped) +def capped(a, b): + return a + b + + +def test_a_field_validator_does_not_fire_on_a_reference_at_the_call(): + """The boundary, pinned: a rule is not judged against a value that does not exist. + + A field written as a link is checked for its shape at the call and for its + rule at the run edge, on what the link delivered. + """ + with Graph(name="linked") as graph: + source = a_thousand() + capped(a=source.result, b=1) + assert links_of(graph) == ["a_thousand.outputs.result -> capped.inputs.a"] + + +def test_the_rule_the_reference_escaped_fires_when_the_link_delivers(): + """The other half of the same claim: 1000 is refused, one checkpoint later.""" + with Graph(name="delivered") as graph: + source = a_thousand() + capped(a=source.result, b=1) + with pytest.raises(TaskInputValidationError, match="at most 100"): + graph.run() + + +def test_the_control_the_same_value_written_as_a_literal_is_refused_at_the_call(): + """What the reference bought: written plainly, the rule answers immediately.""" + with pytest.raises(TaskInputValidationError, match="at most 100"): + with Graph(name="literal"): + capped(a=1000, b=1) + + +def test_a_reference_nested_inside_a_value_leaves_the_whole_field_waiting(): + """A rule reads the field, so one member still to arrive holds the rule back.""" + + class Weights(BaseModel): + weights: dict[str, int] + + @field_validator("weights") + @classmethod + def _positive(cls, value): + if any(item < 0 for item in value.values()): + raise ValueError("every weight must be positive") + return value + + @task(input_model=Weights) + def weighted(weights): + return sum(weights.values()) + + with Graph(name="nested") as graph: + source = a_thousand() + weighted(weights={"a": -1, "b": source.result}) + assert graph.tasks["weighted"] is not None + with pytest.raises(TaskInputValidationError, match="every weight must be positive"): + with Graph(name="whole"): + weighted(weights={"a": -1, "b": 2}) + + +def test_the_private_name_the_rules_are_read_from_still_answers(): + """One canary: pydantic's decorator record is what the rule shadow is built from.""" + from node_graph.input_model import _own_field_validators + + decorators = Capped.__pydantic_decorators__.field_validators + assert set(decorators) == {"cap"} + assert decorators["cap"].info.mode == "after" + assert decorators["cap"].info.fields == ("a",) + assert set(_own_field_validators(Capped)) == {"cap"} + + +def test_a_model_validator_still_waits_for_the_run_edge(): + """A cross-field rule may read a field no one has written yet.""" + + class Ordered(BaseModel): + low: int + high: int + + @model_validator(mode="after") + def _ordered(self): + if self.low >= self.high: + raise ValueError("low must be below high") + return self + + @task(input_model=Ordered) + def ordered(low, high): + return high - low + + graph = Graph(name="unordered") + assert graph.add_task(ordered, "o", low=9, high=1) is not None + with pytest.raises(TaskInputValidationError, match="low must be below high"): + ordered.run(low=9, high=1) + + +#: What each field rule below was handed, in the order the rules ran. +#: +#: The models are written at module level so a task can name them: a task +#: carries a model it cannot import by value, and a copy's rule would write +#: to a copy of this list. +RULE_DATA: list = [] + + +class Bounds(BaseModel): + low: int + high: int + + @field_validator("high") + @classmethod + def _above_low(cls, value, info): + RULE_DATA.append(dict(info.data)) + if value <= info.data["low"]: + raise ValueError("high must exceed low") + return value + + +@task(input_model=Bounds) +def bounds(low, high): + return high - low + + +def test_a_field_rule_reaching_for_a_sibling_waits_until_it_is_there(): + """A rule reading a field nobody wrote cannot be answered on that write. + + What the rule is handed is the discriminating part: the sibling is absent, + so the lookup raises and the rule waits. A placeholder standing in for it + would be judged instead, and judged wrong. + """ + RULE_DATA.clear() + + graph = Graph(name="sibling") + assert graph.add_task(bounds, "b", high=5) is not None + assert RULE_DATA == [{}] + with pytest.raises(TaskInputValidationError, match="high must exceed low"): + bounds.run(low=9, high=5) + assert RULE_DATA[-1] == {"low": 9} + + +class Flagged(BaseModel): + strict: bool = False + value: int = 0 + + @field_validator("value") + @classmethod + def _capped_when_strict(cls, value, info): + RULE_DATA.append(dict(info.data)) + if info.data.get("strict") and value > 10: + raise ValueError("value is capped at 10 in strict mode") + return value + + +@task(input_model=Flagged) +def flagged(strict, value): + return value + + +def test_a_rule_reading_an_unwritten_sibling_reads_it_as_absent(): + """Absent, not defaulted and not a stand-in, so the rule answers as the model does.""" + RULE_DATA.clear() + + # The model's own verdict on the same write, which the wiring check must + # not contradict: the default stands in only where the model itself runs. + assert Flagged(value=99).value == 99 + assert RULE_DATA == [{"strict": False}] + + graph = Graph(name="flagged") + assert graph.add_task(flagged, "f", value=99) is not None + assert RULE_DATA[-1] == {} + with pytest.raises(TaskInputValidationError, match="capped at 10"): + graph.add_task(flagged, "g", strict=True, value=99) + assert RULE_DATA[-1] == {"strict": True} + + +class Inner(BaseModel): + n: int = 1 + + def doubled(self) -> int: + return self.n * 2 + + +class Outer(BaseModel): + inner: Inner = Inner() + + @field_validator("inner") + @classmethod + def _within_reach(cls, value): + if not isinstance(value, Inner): + raise ValueError("inner is not an Inner") + if value.doubled() > 100: + raise ValueError("inner.n is capped at 50") + return value + + +@task(input_model=Outer) +def outer(inner): + return inner + + +def test_a_rule_on_a_nested_field_gets_the_class_the_field_names(): + """It reads the value its field declares, so ``isinstance`` holds and methods answer.""" + graph = Graph(name="nested_ok") + assert graph.add_task(outer, "a", inner={"n": 3}) is not None + assert graph.add_task(outer, "b", inner=Inner(n=3)) is not None + + +def test_and_refuses_the_nested_value_the_model_refuses(): + """The other half: a rule that can read the instance can also judge it.""" + assert Outer(inner={"n": 3}).inner.n == 3 + with pytest.raises(ValidationError, match="capped at 50"): + Outer(inner={"n": 999}) + graph = Graph(name="nested_bad") + with pytest.raises(TaskInputValidationError, match="capped at 50"): + graph.add_task(outer, "c", inner={"n": 999}) + + +def test_a_link_deeper_than_the_walk_goes_leaves_the_field_waiting(): + """Unproved is not resolved: what the walk did not reach counts as awaiting a link.""" + from node_graph.input_model import _MAX_ANNOTATION_DEPTH + + class Deep(BaseModel): + payload: Any = None + + @field_validator("payload") + @classmethod + def _bottoms_out_in_an_int(cls, value): + innermost = value + while isinstance(innermost, list) and innermost: + innermost = innermost[0] + if not isinstance(innermost, int): + raise ValueError("payload must bottom out in an int") + return value + + @task(input_model=Deep) + def deep(payload): + return payload + + graph = Graph(name="deep") + source = graph.add_task(a_thousand, "s") + # One list below where the walk stops, so the socket is never looked at. + buried = source.outputs.result + for _ in range(_MAX_ANNOTATION_DEPTH + 2): + buried = [buried] + assert graph.add_task(deep, "d", payload=buried) is not None + assert graph.add_task(deep, "e", payload=[source.outputs.result]) is not None + # The control: on a value with nothing left to arrive, the rule does run. + with pytest.raises(TaskInputValidationError, match="bottom out in an int"): + graph.add_task(deep, "f", payload=[["not an int"]]) + + +def test_a_rule_on_an_aliased_field_fires_under_either_name(): + """A socket is named by its field, and a write naming the alias reaches it too.""" + + class Priced(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + amount: int = Field(alias="qty") + + @field_validator("amount") + @classmethod + def _not_negative(cls, value): + if value < 0: + raise ValueError("amount cannot be negative") + return value + + @task(input_model=Priced) + def priced(amount): + return amount + + graph = Graph(name="aliased_rule") + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + graph.add_task(priced, "a", amount=-1) + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + graph.add_task(priced, "b", qty=-1) + assert graph.add_task(priced, "c", amount=5) is not None + + +def test_a_rewriting_field_validator_leaves_the_written_value_alone(): + """The rule runs on a copy: what was written is what reaches storage.""" + + class Shouted(BaseModel): + text: str + + @field_validator("text") + @classmethod + def _shout(cls, value): + return value.upper() + + @task(input_model=Shouted) + def shouted(text): + return text + + graph = Graph(name="quiet") + node = graph.add_task(shouted, "s", text="silicon") + assert node.inputs.text.value == "silicon" + + +def test_a_before_validator_is_not_honoured_at_the_call(): + """A normalizer is not a type, so the call sees only what the annotation admits.""" + + class Narrow(BaseModel): + kpoints: list[list[int]] + + @field_validator("kpoints", mode="before") + @classmethod + def as_rows(cls, value): + if value and not isinstance(value[0], list): + return [value] + return value + + @task(input_model=Narrow) + def narrow(kpoints): + return kpoints + + @task.graph() + def flat(m): + return narrow(kpoints=[1, 2, 3]) + + with pytest.raises(TaskInputValidationError, match="kpoints"): + flat.build(m=1) + + +def test_a_widened_annotation_admits_what_the_normalizer_would_accept(): + """The fix for the case above: say in the type what the field takes.""" + + class Widened(BaseModel): + kpoints: list[int] | list[list[int]] + + @task(input_model=Widened) + def widened(kpoints): + return kpoints + + @task.graph() + def flat(m): + return widened(kpoints=[1, 2, 3]) + + assert flat.build(m=1) is not None + + +# -------------------------------------------------------------------------- +# 4. Checkpoint B -- a graph's resolved inputs +# -------------------------------------------------------------------------- + + +class SpanInputs(BaseModel): + """Two ints the socket layer 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.graph(input_model=SpanInputs) +def span(low, high): + return add(x=low, y=high) + + +#: What the graph body below saw, read back after the build. +SEEN: dict[str, str] = {} + + +@task.graph(input_model=SpanInputs) +def watched(low, high): + from node_graph.socket import TaggedValue + + # By kind, not by class name: what a tagged value is spelled as is the + # tagging layer's business, and it has more than one class. + SEEN["low"] = "tagged" if isinstance(low, TaggedValue) else type(low).__name__ + return add(x=low, y=high) + + +def test_a_graph_expands_when_its_inputs_satisfy_the_model(): + assert len(span.build(low=1, high=3).links) == 3 + + +def test_a_cross_field_rule_fires_when_the_graph_is_built(): + with pytest.raises(TaskInputValidationError) as excinfo: + span.build(low=9, high=3) + assert "Graph 'span'" in str(excinfo.value) + assert "low must be below high" in str(excinfo.value) + + +def test_a_field_constraint_the_socket_cannot_express_fires_at_the_same_place(): + with pytest.raises(TaskInputValidationError, match="less than or equal to 100"): + span.build(low=1, high=500) + + +def test_a_cross_field_rule_passes_without_checkpoint_b(monkeypatch): + """The negative control: no other layer knows the two fields are related.""" + import node_graph.utils.graph as graph_utils + + monkeypatch.setattr( + graph_utils, + "_validate_graph_body_inputs", + lambda func, inputs, name: inputs, + ) + assert span.build(low=9, high=3) is not None + + +def test_a_value_of_the_wrong_type_never_reaches_the_body(): + """Where the socket layer's own check does not reach, the model's does.""" + + class Recipe(BaseModel): + amount: Decimal + + @task.graph(input_model=Recipe) + def recipe(amount): + return add(x=1, y=2) + + with pytest.raises(TaskInputValidationError, match="amount"): + recipe.build(amount="not a number") + + +def test_the_body_still_receives_the_tagged_values_it_needs(): + """Validation builds fresh objects; the body is handed the originals, so links survive.""" + SEEN.clear() + graph = watched.build(low=1, high=3) + assert SEEN["low"] == "tagged" + assert "graph_inputs.inputs.low -> add.inputs.x" in links_of(graph) + + +class Doubling(BaseModel): + """``total`` is twice ``count`` -- a rule, stated as a rule.""" + + count: int + total: int + + @model_validator(mode="after") + def check_total(self): + if self.total != self.count * 2: + raise ValueError("total must be twice count") + return self + + +@task.graph(input_model=Doubling) +def doubling(count, total): + return add(x=count, y=total) + + +def test_a_nested_graph_checks_the_value_the_run_produced(): + """A subgraph's inputs are literals only from its own point of view.""" + + @task() + def wrong_total(count): + return count * 3 + + @task.graph() + def outer(count): + computed = wrong_total(count=count) + return doubling(count=count, total=computed.result) + + graph = Graph(name="nested", outputs=ns(result=Any)) + node = graph.add_task(outer, "outer", count=2) + graph.add_link(node.outputs.result, graph.outputs.result) + + with pytest.raises(TaskInputValidationError, match="total must be twice count"): + LocalEngine().run(graph) + + +# -------------------------------------------------------------------------- +# 5. Validation may change how a value is spelled, not what it says +# -------------------------------------------------------------------------- + + +class Derived(BaseModel): + """A model that fills ``total`` in rather than checking it.""" + + count: int + total: int + + @model_validator(mode="after") + def derive_total(self): + object.__setattr__(self, "total", self.count * 2) + return self + + +@task.graph(input_model=Derived) +def derived(count, total): + return add(x=count, y=total) + + +def test_a_derivation_that_leaves_a_resolved_value_alone_is_accepted(): + assert derived.build(count=2, total=4) is not None + + +def test_a_derivation_that_rewrites_its_input_is_refused(): + with pytest.raises(ModelDerivedValueError) as excinfo: + derived.build(count=2, total=7) + assert "'derived'" in str(excinfo.value) + assert "'total'" in str(excinfo.value) + assert "changed it from 7 to 4" in str(excinfo.value) + + +class Coercions(BaseModel): + """Three fields whose stored spelling differs from the object the body wants.""" + + amount: Decimal + color: Color + pair: tuple[int, int] + + +@task.graph(input_model=Coercions) +def coercions(amount, color, pair): + return add(x=1, y=2) + + +def test_pure_coercion_is_not_a_change_of_content(): + assert coercions.build(amount="60", color="red", pair=[1, 2]) is not None + + +def test_a_default_filling_an_omitted_field_is_not_a_change(): + class WithDefault(BaseModel): + x: int + y: int = 7 + + @task.graph(input_model=WithDefault) + def with_default(x, y): + return add(x=x, y=y) + + assert with_default.build(x=1) is not None + + +def test_a_rewriting_validator_is_refused_at_the_run_edge_too(): + class Shouting(BaseModel): + name: str + + @field_validator("name") + @classmethod + def shout(cls, value): + return value.upper() + + @task(input_model=Shouting) + def shouting(name): + return name + + with pytest.raises(ModelDerivedValueError, match="'name'"): + shouting.run(name="silicon") + # Already resolved, so the rule changes nothing and the task runs. + assert shouting.run(name="SILICON") == "SILICON" + + +# -------------------------------------------------------------------------- +# 6. Checkpoint C -- the run edge +# -------------------------------------------------------------------------- + + +class BodyRan(RuntimeError): + """Raised by a body that should not have been reached.""" + + +@task(input_model=SpanInputs) +def announcing_span(low, high): + raise BodyRan + + +@task(inputs=spec_from_model(SpanInputs)) +def announcing_span_without_model(low, high): + raise BodyRan + + +def test_the_body_receives_what_the_model_declares(): + @task(input_model=Coercions, outputs=["amount", "color", "pair"]) + def kinds(amount, color, pair): + return { + "amount": type(amount).__name__, + "color": color is Color.RED, + "pair": type(pair).__name__, + } + + assert kinds.run(amount="60", color="red", pair=[1, 2]) == { + "amount": "Decimal", + "color": True, + "pair": "tuple", + } + + +def test_a_stored_value_the_model_refuses_fails_before_the_body_runs(): + """The body announces itself by raising, and what comes back is the model's refusal.""" + with pytest.raises(TaskInputValidationError, match="low must be below high"): + announcing_span.run(low=9, high=3) + + +def test_without_the_model_the_same_values_reach_the_body(): + """The control: the same sockets, no contract, and the rule is never consulted.""" + with pytest.raises(BodyRan): + announcing_span_without_model.run(low=9, high=3) + + +class SumAndProduct(BaseModel): + total: int + product: int + + +def test_the_output_sockets_come_from_the_output_model(): + @task(output_model=SumAndProduct) + def combine(x, y): + return {"total": x + y, "product": x * y} + + assert set(combine._spec.outputs.fields) == {"total", "product"} + assert combine.run(x=2, y=3) == {"total": 5, "product": 6} + + +def test_a_missing_output_fails_at_the_task_that_produced_it(): + @task(output_model=SumAndProduct) + def forgets_product(x, y): + return {"total": x + y} + + with pytest.raises(TaskOutputValidationError, match="product"): + forgets_product.run(x=2, y=3) + + +def test_a_mistyped_output_fails_at_the_task_that_produced_it(): + @task(output_model=SumAndProduct) + def mistypes_total(x, y): + return {"total": "not a number", "product": x * y} + + with pytest.raises(TaskOutputValidationError, match="total"): + mistypes_total.run(x=2, y=3) + + +# -------------------------------------------------------------------------- +# 7. 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 == "node_graph.int" + assert fields["y"].default == 7 + assert plain_add.run(x=2) == 9 + + +def test_a_model_that_is_not_a_model_is_refused(): + with pytest.raises(ModelContractError, match="must be a pydantic BaseModel"): + + @task(input_model=dict) # type: ignore[arg-type] + def bad(x): + return x + + +def test_the_model_still_validates_when_called_directly(): + """The contract lives on the model, so it holds outside a graph too.""" + with pytest.raises(ValidationError): + SpanInputs(low=9, high=3) + + +# -------------------------------------------------------------------------- +# 8. A value the model has no JSON form for +# -------------------------------------------------------------------------- + + +class Opaque: + """An object with no JSON form at all.""" + + +class HoldsAnOpaqueValue(BaseModel): + payload: Any + n: int + + +@task(input_model=HoldsAnOpaqueValue) +def holds_an_opaque_value(payload, n): + return n + + +def test_a_value_with_no_json_form_does_not_break_the_check(): + """It is compared as the object it is, rather than aborting on a serializer error.""" + assert holds_an_opaque_value.run(payload=Opaque(), n=1) == 1 + + +class ReplacesAnOpaqueValue(BaseModel): + payload: Any + n: int + + @model_validator(mode="after") + def _replace(self): + self.payload = Opaque() + return self + + +@task(input_model=ReplacesAnOpaqueValue) +def replaces_an_opaque_value(payload, n): + return n + + +def test_replacing_a_value_with_no_json_form_is_still_refused(): + with pytest.raises(ModelDerivedValueError, match="'payload'"): + replaces_an_opaque_value.run(payload=Opaque(), n=1) + + +def test_a_value_with_no_json_form_is_stored_as_it_stands(): + payload = Opaque() + assert dump_model_field(HoldsAnOpaqueValue, "payload", payload) is payload + + +# -------------------------------------------------------------------------- +# 9. What a body receives for each kind of field +# -------------------------------------------------------------------------- + + +class Marker: + """A class pydantic can only accept instances of.""" + + +class EveryKind(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + text: str + number: int + fraction: float + flag: bool + mapping: dict + items: list + anything: Any + amount: Decimal + marker: Marker + nested: dict[str, int] + + +@task(input_model=EveryKind) +def every_kind( + text, number, fraction, flag, mapping, items, anything, amount, marker, nested +): + return text + + +def _arrival(name): + return every_kind._spec.inputs.fields[name].meta.extras[BODY_RECEIVES] + + +@pytest.mark.parametrize( + "field", + ["text", "number", "fraction", "flag", "mapping", "items", "amount"], +) +def test_a_field_pydantic_can_rebuild_is_read_as_python(field): + """The engine's wrapper comes off, because the model builds the value back.""" + assert _arrival(field) == "python" + + +@pytest.mark.parametrize("field", ["anything", "marker"]) +def test_a_field_only_an_instance_satisfies_is_read_as_a_node(field): + """``Any`` declares nothing to rebuild, and an arbitrary class wants its own instance.""" + assert _arrival(field) == "node" + + +def test_the_member_of_a_typed_mapping_carries_the_mark_too(): + item = every_kind._spec.inputs.fields["nested"].item + assert item.meta.extras[BODY_RECEIVES] == "python" + + +def test_a_nested_models_leaves_carry_the_mark_and_the_namespace_does_not(): + class Inner(BaseModel): + x: int + + class Outer(BaseModel): + inner: Inner + + @task(input_model=Outer) + def outer(inner): + return inner + + namespace = outer._spec.inputs.fields["inner"] + assert BODY_RECEIVES not in namespace.meta.extras + assert namespace.fields["x"].meta.extras[BODY_RECEIVES] == "python" + + +def _kind_of(annotation): + """Return the mark a one-field model gives ``annotation``.""" + + class OneField(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + field: annotation + + @task(input_model=OneField) + def one_field(field): + return field + + return one_field._spec.inputs.fields["field"].meta.extras[BODY_RECEIVES] + + +@pytest.mark.parametrize( + "annotation, expected", + [ + (Optional[int], "python"), + (Optional[Marker], "node"), + (Union[int, str], "python"), + (Union[Marker, Opaque], "node"), + ], + ids=["optional-python", "optional-node", "two-python-arms", "two-node-arms"], +) +def test_a_union_whose_arms_agree_takes_their_answer(annotation, expected): + """``None`` is neither kind, so ``X | None`` is read as ``X``.""" + assert _kind_of(annotation) == expected + + +def test_a_union_spanning_both_kinds_is_refused_at_decoration(): + """A socket delivers one form, so nothing here could choose between them.""" + with pytest.raises(ModelContractError) as excinfo: + _kind_of(Union[int, Marker]) + message = str(excinfo.value) + assert "OneField.field" in message + assert "int" in message and "Marker" in message + + +def test_a_none_arm_does_not_make_a_mixed_union_agree(): + """The control: dropping ``None`` leaves the two disagreeing arms it was hiding.""" + with pytest.raises(ModelContractError, match="OneField.field"): + _kind_of(Optional[Union[int, Marker]]) + + +# -------------------------------------------------------------------------- +# 10. The model's own config decides what its fields accept +# -------------------------------------------------------------------------- + + +class Stripped(BaseModel): + model_config = ConfigDict(str_strip_whitespace=True) + + text: str + + +@task(input_model=Stripped) +def stripped(text): + return text + + +def test_a_config_that_coerces_is_not_a_derivation(): + """Both sides strip, because both carry the config, so the content is unchanged.""" + assert stripped.run(text=" silicon ") == "silicon" + + +class StripsByHand(BaseModel): + text: str + + @field_validator("text") + @classmethod + def _strip(cls, value): + return value.strip() + + +@task(input_model=StripsByHand) +def strips_by_hand(text): + return text + + +def test_the_control_a_validator_doing_the_same_thing_is_a_derivation(): + """A rule is the model's own; only one side runs it, so the content changes.""" + with pytest.raises(ModelDerivedValueError, match="'text'"): + strips_by_hand.run(text=" silicon ") + + +class Aliased(BaseModel): + model_config = ConfigDict(populate_by_name=True) + + amount: int = Field(alias="qty") + + +@task(input_model=Aliased) +def aliased(amount): + return amount + + +def test_a_field_reachable_by_its_own_name_is_accepted_at_the_call(): + assert aliased.run(amount=5) == 5 + + +def test_an_arbitrary_type_is_accepted_where_its_config_allows_it(): + marker = Marker() + assert ( + every_kind.run( + text="a", + number=1, + fraction=1.0, + flag=True, + mapping={}, + items=[], + anything="whatever", + amount=Decimal("1"), + marker=marker, + nested={"k": 1}, + ) + == "a" + ) + + +# -------------------------------------------------------------------------- +# 11. A field kept out of the model's dump +# -------------------------------------------------------------------------- + + +class Excluded(BaseModel): + secret: str = Field(exclude=True) + n: int + + @field_validator("secret") + @classmethod + def _rewrite(cls, value): + return value + "-REWRITTEN" + + +@task(input_model=Excluded) +def excluded(secret, n): + return secret + + +def test_a_field_kept_out_of_a_dump_is_still_compared(): + """``exclude=True`` says what is rendered, not what may be rewritten unseen.""" + with pytest.raises(ModelDerivedValueError, match="'secret'"): + excluded.run(secret="abc", n=1) + + +class HiddenByItsAnnotation(BaseModel): + """Every value renders as ``0``, so a rewrite leaves no trace in the dump.""" + + amount: Annotated[int, PlainSerializer(lambda value: 0, return_type=int)] + + @field_validator("amount") + @classmethod + def _double(cls, value): + return value * 2 + + +@task(input_model=HiddenByItsAnnotation) +def hidden_by_its_annotation(amount): + return amount + + +def test_a_serializer_written_into_the_annotation_cannot_hide_a_rewrite(): + """It renders, so the twin drops it and compares what the field actually says.""" + with pytest.raises(ModelDerivedValueError, match="changed it from 1 to 2"): + hidden_by_its_annotation.run(amount=1) + + +def test_the_same_serializer_still_decides_the_stored_form(): + """The control: dropping it from the twin does not drop it from storage.""" + assert dump_model_field(HiddenByItsAnnotation, "amount", 7) == 0 + + +def test_an_excluded_fields_default_is_stored_as_it_stands(): + """``dump_model_field`` has no rendered form to give, and says the value instead.""" + + class WithExcludedDefault(BaseModel): + secret: str = Field(default="abc", exclude=True) + + assert dump_model_field(WithExcludedDefault, "secret", "abc") == "abc" + + +# -------------------------------------------------------------------------- +# 12. The reference the invariance check compares against +# -------------------------------------------------------------------------- + + +class MutatesInPlace(BaseModel): + rows: Any + n: int + + @model_validator(mode="after") + def _append(self): + self.rows.append("INJECTED") + return self + + +@task(input_model=MutatesInPlace) +def mutates_in_place(rows, n): + return list(rows) + + +def test_a_validator_that_rewrites_its_input_in_place_is_refused(): + """The reference is read before the model runs, so there is still one to compare with.""" + with pytest.raises(ModelDerivedValueError, match="'rows'"): + mutates_in_place.run(rows=["a"], n=1) + + +class Rebinds(BaseModel): + rows: list + n: int + + @model_validator(mode="after") + def _rebind(self): + self.rows = [*self.rows, "INJECTED"] + return self + + +@task(input_model=Rebinds) +def rebinds(rows, n): + return list(rows) + + +def test_a_validator_that_replaces_its_input_is_refused(): + with pytest.raises(ModelDerivedValueError, match="'rows'"): + rebinds.run(rows=["a"], n=1) + + +class Leaf(BaseModel): + x: int + + +class ReachesInside(BaseModel): + model_config = ConfigDict(revalidate_instances="never") + + leaf: Leaf + n: int + + @model_validator(mode="after") + def _reach(self): + self.leaf.x = 999 + return self + + +@task(input_model=ReachesInside) +def reaches_inside(leaf, n): + return leaf.x + + +def test_a_validator_that_reaches_into_a_nested_instance_is_refused(): + """The caller's own instance is mutated, so only a reference taken first sees it.""" + with pytest.raises(ModelDerivedValueError, match="'leaf'"): + reaches_inside.run(leaf=Leaf(x=1), n=1) + + +def test_without_taking_the_reference_first_the_mutation_is_invisible(): + """The control: read the reference after the model ran and both sides say the same.""" + from node_graph import input_model as module + + given = {"rows": ["a"], "n": 1} + validated = MutatesInPlace.model_validate(dict(given)) + after_the_fact = module.content_snapshot(MutatesInPlace, given) + module.check_content_invariance( + MutatesInPlace, after_the_fact, validated, label="control" + ) + assert given["rows"] == ["a", "INJECTED"] + + +# -------------------------------------------------------------------------- +# 13. An annotation's own validators run on both sides +# -------------------------------------------------------------------------- + + +class AnnotatedIdempotent(BaseModel): + amount: Annotated[int, AfterValidator(lambda value: abs(value))] + + +@task(input_model=AnnotatedIdempotent) +def annotated_idempotent(amount): + return amount + + +def test_an_idempotent_annotation_validator_passes(): + """It runs on both sides and lands on the same value, so nothing changed.""" + assert annotated_idempotent.run(amount=-6) == 6 + + +class AnnotatedDoubling(BaseModel): + amount: Annotated[int, AfterValidator(lambda value: value * 2)] + + +@task(input_model=AnnotatedDoubling) +def annotated_doubling(amount): + return amount + + +def test_a_non_idempotent_annotation_validator_is_refused(): + """Running twice doubles twice, so the two sides disagree and the write is refused.""" + with pytest.raises(ModelDerivedValueError, match="changed it from 6 to 12"): + annotated_doubling.run(amount=3) + + +# -------------------------------------------------------------------------- +# 14. A nested field accepts the class it names +# -------------------------------------------------------------------------- + + +class Point(BaseModel): + x: int + + +class HoldsAPoint(BaseModel): + point: Point + + +class HoldsPointList(BaseModel): + points: list[Point] + + +class HoldsPointMapping(BaseModel): + points: dict[str, Point] + + +@task(input_model=HoldsAPoint) +def holds_a_point(point): + return point + + +@task(input_model=HoldsPointList) +def holds_point_list(points): + return points + + +@task(input_model=HoldsPointMapping) +def holds_point_mapping(points): + return points + + +@pytest.mark.parametrize( + "handle, written", + [ + (holds_a_point, {"point": Point(x=1)}), + (holds_point_list, {"points": [Point(x=1)]}), + (holds_point_mapping, {"points": {"k": Point(x=1)}}), + ], + ids=["plain", "in-a-list", "in-a-mapping"], +) +def test_the_class_a_field_names_is_accepted_where_it_is_written(handle, written): + """A shadow is a different class; the class the field names is not a wrong value.""" + with Graph(name="instances"): + assert handle(**written) is not None + + +@pytest.mark.parametrize( + "handle, written", + [ + (holds_a_point, {"point": {"x": "nope"}}), + (holds_point_list, {"points": [{"x": "nope"}]}), + (holds_point_mapping, {"points": {"k": {"x": "nope"}}}), + ], + ids=["plain", "in-a-list", "in-a-mapping"], +) +def test_a_wrong_value_in_the_same_place_is_still_refused(handle, written): + """The control: accepting the class does not stop the shadow checking a dict.""" + with Graph(name="instances_bad"): + with pytest.raises(TaskInputValidationError, match="x"): + handle(**written) + + +# -------------------------------------------------------------------------- +# 15. Checkpoint A fires wherever an input is written +# -------------------------------------------------------------------------- + + +class Bounded(BaseModel): + amount: int = Field(gt=0) + other: int = 5 + + +@task(input_model=Bounded) +def bounded(amount, other): + return amount + + +def test_a_bad_value_is_refused_when_the_task_is_called(): + with pytest.raises(TaskInputValidationError, match="amount"): + with Graph(name="call"): + bounded(amount="bad", other=1) + + +def test_a_bad_value_is_refused_when_it_is_passed_to_add_task(): + graph = Graph(name="added") + with pytest.raises(TaskInputValidationError, match="amount"): + graph.add_task(bounded, "b", amount="bad") + + +def test_a_bad_value_is_refused_when_it_is_set_on_the_task(): + graph = Graph(name="set") + node = graph.add_task(bounded, "b") + with pytest.raises(TaskInputValidationError, match="amount"): + node.set_inputs({"amount": "bad"}) + + +def test_a_constraint_the_socket_cannot_see_is_refused_at_add_task(): + graph = Graph(name="constrained") + with pytest.raises(TaskInputValidationError, match="greater than 0"): + graph.add_task(bounded, "b", amount=-1) + + +def test_without_the_check_add_task_accepts_the_constraint_breach(monkeypatch): + """The control: the socket layer types ``amount`` as ``int`` and cannot see ``gt=0``.""" + from node_graph import input_model as module + + monkeypatch.setattr(module, "validate_task_inputs", lambda task, inputs: None) + graph = Graph(name="unchecked") + assert graph.add_task(bounded, "b", amount=-1) is not None + + +def test_assigning_to_a_socket_writes_past_the_check(): + """The one write route the model does not reach, pinned so the claim cannot rot. + + The socket layer types ``amount`` as ``int`` and carries no ``gt=0``, so + the assignment is refused by nothing and the value is stored; + ``set_inputs`` refuses the same value. + """ + graph = Graph(name="assigned") + node = graph.add_task(bounded, "b") + node.inputs.amount.value = -1 + assert node.inputs.amount.value == -1 + with pytest.raises(TaskInputValidationError, match="greater than 0"): + node.set_inputs({"amount": -1}) + + +def test_writing_only_some_of_the_fields_is_not_a_missing_input(): + """Inputs may be written a few at a time, and a link may supply the rest.""" + graph = Graph(name="partial") + assert graph.add_task(bounded, "b") is not None + assert graph.add_task(bounded, "c", amount=1) is not None + + +def test_a_task_written_into_an_input_is_a_reference_not_a_value(): + graph = Graph(name="linked") + producer = graph.add_task(bounded, "p", amount=1) + assert graph.add_task(bounded, "c", amount=producer.outputs.result) is not None + + +class Rated(BaseModel): + """The rule is the model's own, so no socket type could carry it.""" + + rating: int = 1 + + @field_validator("rating") + @classmethod + def _in_range(cls, value): + if not 1 <= value <= 5: + raise ValueError("rating runs from 1 to 5") + return value + + +@task(input_model=Rated) +def rated(rating): + return rating + + +def test_a_field_rule_is_refused_when_it_is_passed_to_add_task(): + graph = Graph(name="rated_added") + with pytest.raises(TaskInputValidationError, match="rating runs from 1 to 5"): + graph.add_task(rated, "r", rating=9) + + +def test_a_field_rule_is_refused_when_it_is_set_on_the_task(): + graph = Graph(name="rated_set") + node = graph.add_task(rated, "r") + with pytest.raises(TaskInputValidationError, match="rating runs from 1 to 5"): + node.set_inputs({"rating": 9}) + + +def test_without_the_check_add_task_accepts_the_broken_rule(without_wiring_checks): + """The control: the socket layer types ``rating`` as ``int`` and sees no rule.""" + graph = Graph(name="rated_unchecked") + assert graph.add_task(rated, "r", rating=9) is not None + + +# -------------------------------------------------------------------------- +# 16. A value whose inequality is not a yes or no +# -------------------------------------------------------------------------- + + +class HoldsAnArray(BaseModel): + payload: Any + n: int + + +@task(input_model=HoldsAnArray) +def holds_an_array(payload, n): + return n + + +class ReplacesTheArray(BaseModel): + payload: Any + n: int + + @model_validator(mode="after") + def _replace(self): + import numpy + + self.payload = numpy.array([9, 9]) + return self + + +@task(input_model=ReplacesTheArray) +def replaces_the_array(payload, n): + return n + + +@pytest.mark.parametrize("size", [1, 2], ids=["one-element", "two-elements"]) +def test_an_array_the_model_leaves_alone_passes(size): + """``array != array`` answers with an array, which the check must not read as a bool.""" + numpy = pytest.importorskip("numpy") + assert holds_an_array.run(payload=numpy.arange(size), n=1) == 1 + + +@pytest.mark.parametrize("size", [1, 2], ids=["one-element", "two-elements"]) +def test_replacing_an_array_is_still_refused(size): + """The control: the comparison still decides, it just cannot ask ``if``.""" + numpy = pytest.importorskip("numpy") + with pytest.raises(ModelDerivedValueError, match="'payload'"): + replaces_the_array.run(payload=numpy.arange(size), n=1) + + +# -------------------------------------------------------------------------- +# 17. A container of values only an instance satisfies +# -------------------------------------------------------------------------- + + +def _spec_of(annotation): + """Return the spec a one-field model gives ``annotation``.""" + + class OneField(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + field: annotation + + return spec_from_model(OneField).fields["field"] + + +class MarkerRow(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + marker: Marker + + +class MixedRow(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + marker: Marker + count: int + + +class MarkerEntry(TypedDict): + marker: Marker + + +class MixedEntry(TypedDict): + marker: Marker + count: int + + +def test_a_list_of_values_only_an_instance_satisfies_is_refused(): + """Storage keeps a list as one node of plain data, so the members cannot come back.""" + with pytest.raises(ModelContractError, match="cannot round-trip through storage"): + _spec_of(list[Marker]) + + +def test_the_refusal_names_the_mapping_that_does_work(): + with pytest.raises(ModelContractError, match=r"dict\[str, Marker\]"): + _spec_of(list[Marker]) + + +def test_a_mapping_of_the_same_type_is_the_shape_that_survives(): + """The control: one socket, and one node, per key.""" + spec = _spec_of(dict[str, Marker]) + assert spec.meta.dynamic is True + assert spec.item.meta.extras[BODY_RECEIVES] == "node" + + +def test_a_list_of_plain_data_is_untouched(): + """The refusal is about what the members are, not about the list.""" + assert _spec_of(list[int]).meta.extras[BODY_RECEIVES] == "python" + + +def test_a_list_of_anything_is_untouched(): + """``Any`` declares nothing, so it declares nothing to refuse either.""" + assert _spec_of(list[Any]).meta.extras[BODY_RECEIVES] == "python" + + +def test_a_typed_dict_of_engine_types_arrives_as_the_engine_stored_it(): + """Pydantic builds the mapping readily; what is inside it still wants an instance.""" + assert _spec_of(MarkerEntry).meta.extras[BODY_RECEIVES] == "node" + + +def test_a_typed_dict_whose_members_disagree_is_refused(): + with pytest.raises(ModelContractError, match="whose members disagree"): + _spec_of(MixedEntry) + + +def test_a_model_in_a_list_is_read_through_its_own_fields(): + """``list[MarkerRow]`` is a list of values only an instance satisfies, once removed.""" + with pytest.raises(ModelContractError, match="cannot round-trip through storage"): + _spec_of(list[MarkerRow]) + + +def test_a_model_in_a_list_whose_members_disagree_is_refused(): + with pytest.raises(ModelContractError, match="whose members disagree"): + _spec_of(list[MixedRow]) + + +def test_a_list_of_models_of_plain_data_still_works(): + """The control: the same shape, with members the model can rebuild.""" + assert _spec_of(list[Point]).meta.extras[BODY_RECEIVES] == "python" + + +def test_a_union_refusal_reads_as_a_sentence(): + """One arm on each side, so the verbs are singular and ``Any`` says what it is.""" + with pytest.raises(ModelContractError) as excinfo: + _spec_of(Union[int, Marker]) + message = str(excinfo.value) + assert "int is rebuilt from plain data" in message + assert "Marker arrives as the engine stored it" in message + + +def test_an_any_arm_is_named_for_what_it_declares(): + with pytest.raises(ModelContractError, match="Any declares nothing to rebuild"): + _spec_of(Union[int, Any]) + + +# -------------------------------------------------------------------------- +# 18. A field reachable only by its alias +# -------------------------------------------------------------------------- + + +class AliasedOnly(BaseModel): + """No ``populate_by_name``: as written, the model takes ``qty`` and not ``amount``.""" + + amount: int = Field(alias="qty") + + +@task(input_model=AliasedOnly) +def aliased_only(amount): + return amount + + +def test_the_socket_is_named_by_the_field_not_the_alias(): + assert set(aliased_only._spec.inputs.fields) == {"amount"} + + +def test_an_aliased_field_is_accepted_under_the_name_that_names_its_socket(): + """The socket delivers ``amount``, so every checkpoint has to take it.""" + assert aliased_only.run(amount=5) == 5 + + +def test_an_aliased_field_is_still_accepted_under_its_alias(): + """Widening what is accepted takes nothing away from the model as written.""" + assert AliasedOnly.model_validate({"qty": 5}).amount == 5 + + +def test_a_field_renamed_by_an_alias_generator_is_accepted_too(): + class Generated(BaseModel): + model_config = ConfigDict(alias_generator=lambda name: name.upper()) + + amount: int + + @task(input_model=Generated) + def generated(amount): + return amount + + assert generated.run(amount=5) == 5 + + +def test_the_model_as_written_still_refuses_the_field_name(): + """The control: the acceptance is added for the contract, not taken from the user.""" + with pytest.raises(ValidationError): + AliasedOnly.model_validate({"amount": 5}) + + +# -------------------------------------------------------------------------- +# 19. What the body returns +# -------------------------------------------------------------------------- + + +class OpaqueOutput(BaseModel): + model_config = ConfigDict(arbitrary_types_allowed=True) + + payload: Marker + n: int + + +def test_an_output_with_no_json_form_is_returned_as_it_stands(): + """One field the model cannot render does not cost the others theirs. + + Dumping the whole model raises on ``Marker``; per field, only that field + has no rendering and it keeps the object instead. The engine copies the + object on its way out, so it is compared by type. + """ + + @task(output_model=OpaqueOutput) + def emits(x): + return {"payload": Marker(), "n": x} + + returned = emits.run(x=3) + assert isinstance(returned["payload"], Marker) + assert returned["n"] == 3 + + +def test_a_returned_key_the_model_does_not_declare_is_refused(): + """It has no socket to be written to, so it would be dropped between here and there.""" + + @task(output_model=SumAndProduct) + def emits_extra(x, y): + return {"total": x + y, "product": x * y, "difference": x - y} + + with pytest.raises(TaskOutputValidationError, match="'difference'"): + emits_extra.run(x=3, y=2) + + +def test_the_declared_keys_alone_still_pass(): + """The control: the same body without the extra key is accepted.""" + + @task(output_model=SumAndProduct) + def emits_declared(x, y): + return {"total": x + y, "product": x * y} + + assert emits_declared.run(x=3, y=2) == {"total": 5, "product": 6} + + +# -------------------------------------------------------------------------- +# 20. A tag is judged by the value it carries +# -------------------------------------------------------------------------- + + +class Cutoff(BaseModel): + """A constraint and no rule, so only the type pass can enforce it.""" + + structure: str + ecutwfc: float = Field(default=60.0, le=200) + + +@task(input_model=Cutoff) +def cutoff(structure, ecutwfc): + return ecutwfc + + +class CutoffCaller(BaseModel): + structure: str = "si" + ecutwfc: float = 60.0 + + +@task.graph(input_model=CutoffCaller) +def calls_cutoff(structure, ecutwfc): + return cutoff(structure=structure, ecutwfc=ecutwfc).result + + +def test_a_constraint_is_enforced_on_a_value_a_graph_body_passes_on(): + """A graph body hands its inputs on as tags, and a tag carries a value. + + The graph itself admits any cutoff; the bound is the leaf's own, and it is + answered at the line inside the body that writes the value. + """ + with pytest.raises(TaskInputValidationError, match="less than or equal to 200"): + calls_cutoff.build(structure="si", ecutwfc=500) + + +def test_the_same_body_still_draws_the_link_for_a_value_the_bound_admits(): + """The control: judging the tag costs neither the build nor the link. + + A tag stripped to its value would leave the leaf holding a copy, so the + link is what says the tag survived the check. + """ + graph = calls_cutoff.build(structure="si", ecutwfc=100) + assert "graph_inputs.inputs.ecutwfc -> cutoff.inputs.ecutwfc" in links_of(graph) + + +def test_a_tagged_value_and_the_value_it_carries_are_judged_alike(): + """The unit behind both: the type pass reads through the tag as the rules do.""" + from node_graph.socket import TaggedValue + + with Graph(name="tagged") as graph: + source = a_thousand() + + tagged = TaggedValue(500, socket=source.result) + for written in (500, tagged): + with pytest.raises(TaskInputValidationError, match="less than or equal to 200"): + validate_wiring_inputs( + Cutoff, {"ecutwfc": written}, label="cutoff", complete=False + ) + assert graph is not None + + +def test_a_bare_socket_in_the_same_place_is_still_a_reference(): + """The other half: what a link will deliver is not a value to judge.""" + with Graph(name="reference") as graph: + source = a_thousand() + cutoff(structure="si", ecutwfc=source.result) + assert links_of(graph) == ["a_thousand.outputs.result -> cutoff.inputs.ecutwfc"] + + +# -------------------------------------------------------------------------- +# 21. One rule waiting does not silence the rest of the write +# -------------------------------------------------------------------------- + + +class Job(BaseModel): + """Two rules, one of which reads a sibling.""" + + low: int = 0 + high: int = 10 + cores: int = 1 + + @field_validator("high") + @classmethod + def _above_low(cls, value, info): + if value <= info.data["low"]: + raise ValueError("high must exceed low") + return value + + @field_validator("cores") + @classmethod + def _at_least_one(cls, value): + if value < 1: + raise ValueError("cores must be at least 1") + return value + + +@task(input_model=Job) +def job(low, high, cores): + return cores + + +def test_a_rule_waiting_for_a_sibling_leaves_its_neighbours_running(): + """``high``'s rule has no ``low`` to read; ``cores``'s rule is answerable.""" + graph = Graph(name="job") + with pytest.raises(TaskInputValidationError, match="cores must be at least 1"): + graph.add_task(job, "j", high=5, cores=-4) + + +def test_the_same_two_rules_still_answer_when_the_sibling_is_there(): + """The control: nothing waits, and the write is refused for the same reason.""" + graph = Graph(name="job_whole") + with pytest.raises(TaskInputValidationError, match="cores must be at least 1"): + graph.add_task(job, "j", low=0, high=5, cores=-4) + + +def test_the_rule_that_waited_is_answered_at_the_run_edge(): + """The other half: what waits is judged once the whole payload is in hand.""" + with pytest.raises(TaskInputValidationError, match="high must exceed low"): + job.run(low=9, high=5, cores=1) + + +#: The table the rule below reads, so a missing key is the rule's own error. +KNOWN_KEYS = {"a": 1} + + +class Keyed(BaseModel): + key: str = "a" + other: int = 0 + + @field_validator("key") + @classmethod + def _known(cls, value): + KNOWN_KEYS[value] + return value + + @field_validator("other") + @classmethod + def _positive(cls, value): + if value < 0: + raise ValueError("other must be positive") + return value + + +@task(input_model=Keyed) +def keyed(key, other): + return key + + +def test_a_key_error_of_the_rules_own_making_refuses_the_write(): + """A rule that raises is a rule that answered, whatever it raised. + + The model itself raises ``KeyError`` on the same value, so a write the + wiring check waves through would be a write the model refuses. + """ + with pytest.raises(KeyError): + Keyed(key="zzz") + graph = Graph(name="keyed") + with pytest.raises(KeyError): + graph.add_task(keyed, "k", key="zzz") + + +def test_that_rule_does_not_silence_the_write_it_shares(): + """The control: the neighbouring rule is reached on a key the table has.""" + graph = Graph(name="keyed_ok") + with pytest.raises(TaskInputValidationError, match="other must be positive"): + graph.add_task(keyed, "k", key="a", other=-1) + + +# -------------------------------------------------------------------------- +# 22. A nested model's cross-field rule waits too +# -------------------------------------------------------------------------- + + +class Agreeing(BaseModel): + n: int = 1 + m: int = 1 + + @model_validator(mode="after") + def _agree(self): + if self.n != self.m: + raise ValueError("n and m must agree") + return self + + +class HoldsAgreeing(BaseModel): + """No rule of its own, so nothing of this model's runs at a write.""" + + pair: Agreeing = Agreeing() + note: str = "x" + + +class HoldsAgreeingAndARule(BaseModel): + """The same shape with one field rule, unrelated to the nested model.""" + + pair: Agreeing = Agreeing() + note: str = "x" + + @field_validator("note") + @classmethod + def _short(cls, value): + if len(value) > 10: + raise ValueError("note is capped at 10") + return value + + +@task(input_model=HoldsAgreeing) +def holds_agreeing(pair, note): + return note + + +@task(input_model=HoldsAgreeingAndARule) +def holds_agreeing_and_a_rule(pair, note): + return note + + +@pytest.mark.parametrize( + "handle", [holds_agreeing, holds_agreeing_and_a_rule], ids=["bare", "ruled"] +) +def test_a_nested_cross_field_rule_waits_whatever_else_the_model_carries(handle): + """A write is a write: what the outer model happens to declare cannot decide it. + + The nested rule may read a field a later write or a link supplies, so it + waits exactly as the outer model's own cross-field rules do. + """ + graph = Graph(name=f"nested_{handle.identifier}") + assert graph.add_task(handle, "h", pair={"n": 1, "m": 2}) is not None + + +@pytest.mark.parametrize( + "handle", [holds_agreeing, holds_agreeing_and_a_rule], ids=["bare", "ruled"] +) +def test_the_same_nested_rule_is_answered_at_the_run_edge(handle): + """The other half: the rule is enforced, one checkpoint later, in both models.""" + with pytest.raises(TaskInputValidationError, match="n and m must agree"): + handle.run(pair={"n": 1, "m": 2}, note="x") + + +def test_a_field_rule_inside_a_nested_model_still_runs_at_the_write(): + """What waits is the cross-field rule; a rule on one field is answerable.""" + + class Capped(BaseModel): + n: int = 1 + + @field_validator("n") + @classmethod + def _small(cls, value): + if value > 100: + raise ValueError("n is capped at 100") + return value + + class HoldsCapped(BaseModel): + inner: Capped = Capped() + + @task(input_model=HoldsCapped) + def holds_capped(inner): + return inner + + graph = Graph(name="nested_field_rule") + with pytest.raises(TaskInputValidationError, match="n is capped at 100"): + graph.add_task(holds_capped, "h", inner={"n": 999}) + assert graph.add_task(holds_capped, "ok", inner={"n": 3}) is not None + + +def test_the_private_name_a_models_cross_field_rules_are_read_from_still_answers(): + """One canary: the twin is the model with that record emptied, and only there.""" + from node_graph.input_model import _field_rules_only + + assert set(Agreeing.__pydantic_decorators__.model_validators) == {"_agree"} + twin = _field_rules_only(Agreeing) + assert issubclass(twin, Agreeing) + assert twin.__pydantic_decorators__.model_validators == {} + assert set(Agreeing.__pydantic_decorators__.model_validators) == {"_agree"} + + +# -------------------------------------------------------------------------- +# 23. A model this process cannot read is not a task without rules +# -------------------------------------------------------------------------- + + +class _RaisingExecutor: + """An executor whose callable cannot be produced.""" + + def __init__(self, error: Exception) -> None: + self._error = error + + @property + def callable(self): + raise self._error + + +class _SpecWith: + def __init__(self, executor) -> None: + self.executor = executor + + +class _TaskWith: + def __init__(self, executor) -> None: + self.name = "t" + self.spec = _SpecWith(executor) + + +def test_a_callable_this_process_cannot_rebuild_is_refused_by_name(): + """Reading the model is how the checks are found; failing to read it is loud.""" + from node_graph.input_model import _input_model_of_task + + task_like = _TaskWith(_RaisingExecutor(RuntimeError("no schema"))) + with pytest.raises(ModelContractError, match="importable module"): + _input_model_of_task(task_like) + + +def test_a_callable_that_is_simply_not_here_still_has_no_contract(): + """The control, and the documented case: nothing to import is nothing to enforce.""" + from node_graph.input_model import _input_model_of_task + + task_like = _TaskWith(_RaisingExecutor(ImportError("no such module"))) + assert _input_model_of_task(task_like) is None + + +def test_a_task_whose_callable_is_here_reads_its_model(): + """The other control: the ordinary path still finds the model it enforces.""" + from node_graph.input_model import _input_model_of_task + + graph = Graph(name="resolves") + assert _input_model_of_task(graph.add_task(bounded, "b")) is Bounded + + +def test_a_model_written_in_a_script_that_cannot_be_rebuilt_is_refused(tmp_path): + """The route this reaches in practice, run as a script rather than described. + + Under ``from __future__ import annotations``, a model declaring a nested + model that carries a ``@model_validator`` does not survive the round trip + a ``__main__`` callable is stored through, and a model that cannot be read + is a task whose every write would go unchecked. + """ + import subprocess + import sys + + script = tmp_path / "run.py" + script.write_text( + "from __future__ import annotations\n" + "from pydantic import BaseModel, model_validator\n" + "from node_graph import Graph, task\n" + "\n" + "class Pair(BaseModel):\n" + " n: int = 1\n" + " m: int = 1\n" + " @model_validator(mode='after')\n" + " def _agree(self):\n" + " if self.n != self.m:\n" + " raise ValueError('n and m must agree')\n" + " return self\n" + "\n" + "class Holder(BaseModel):\n" + " pair: Pair = Pair()\n" + "\n" + "@task(input_model=Holder)\n" + "def holder(pair):\n" + " return pair\n" + "\n" + "Graph(name='s').add_task(holder, 'h')\n" + ) + finished = subprocess.run( + [sys.executable, str(script)], capture_output=True, text=True + ) + assert finished.returncode != 0 + assert "ModelContractError" in finished.stderr + assert "importable module" in finished.stderr + + +# -------------------------------------------------------------------------- +# 24. A socket refusal and a model refusal address the same value alike +# -------------------------------------------------------------------------- + + +class Control(BaseModel): + model_config = ConfigDict(extra="forbid") + + calculation: Literal["nscf"] = "nscf" + + +class Parameters(BaseModel): + CONTROL: Control = Control() + + +class Forced(BaseModel): + """A namelist whose one keyword the route decides, so no field declares it.""" + + model_config = ConfigDict(extra="forbid") + + +class Overrides(BaseModel): + nscf: Forced = Forced() + + +class RouteInputs(BaseModel): + parameters: Parameters = Parameters() + overrides: Overrides = Overrides() + + +@task(input_model=RouteInputs) +def route_leaf(parameters, overrides): + return 1 + + +@task.graph(input_model=RouteInputs) +def route(parameters, overrides): + return route_leaf(parameters=parameters, overrides=overrides).result + + +@pytest.mark.parametrize( + "written, loc, error_type", + [ + ( + {"parameters": {"CONTROL": {"calculation": "scf"}}}, + ("parameters", "CONTROL", "calculation"), + "literal_error", + ), + ( + {"overrides": {"nscf": {"nosym": False}}}, + ("overrides", "nscf", "nosym"), + "extra_forbidden", + ), + ], + ids=["outside the literal", "no field to write to"], +) +def test_a_graphs_own_input_is_refused_with_the_path_that_was_refused( + written, loc, error_type +): + """A graph's own inputs are written into sockets, and a socket says where. + + Written into the graph, the value never reaches a model: the refusal comes + from the socket layer, which now carries the same two facts a model's + refusal carries. + """ + from node_graph.errors import SocketValueError + + with pytest.raises(SocketValueError) as caught: + route.build(**written) + assert caught.value.loc == loc + assert caught.value.type == error_type + + +@pytest.mark.parametrize( + "written, loc", + [ + ( + {"parameters": {"CONTROL": {"calculation": "scf"}}}, + ("parameters", "CONTROL", "calculation"), + ), + ( + {"overrides": {"nscf": {"nosym": False}}}, + ("overrides", "nscf", "nosym"), + ), + ], + ids=["outside the literal", "no field to write to"], +) +def test_the_model_refuses_the_same_write_at_the_same_path(written, loc): + """The reference: what the leaf's model says about the same two writes.""" + graph = Graph(name="model_side") + with pytest.raises(TaskInputValidationError) as caught: + graph.add_task(route_leaf, "leaf", **written) + assert [error["loc"] for error in caught.value.__cause__.errors()] == [loc] + + +@pytest.mark.parametrize( + "written, loc, error_type", + [ + ( + {"parameters": {"CONTROL": {"calculation": "scf"}}}, + ("parameters", "CONTROL", "calculation"), + "literal_error", + ), + ( + {"overrides": {"nscf": {"nosym": False}}}, + ("overrides", "nscf", "nosym"), + "extra_forbidden", + ), + ], + ids=["outside the literal", "no field to write to"], +) +def test_a_models_refusal_is_read_the_way_a_sockets_is(written, loc, error_type): + """Where and why, off the refusal itself: neither caller parses a message.""" + graph = Graph(name="model_side_read") + with pytest.raises(TaskInputValidationError) as caught: + graph.add_task(route_leaf, "leaf", **written) + assert [(error["loc"], error["type"]) for error in caught.value.errors] == [ + (loc, error_type) + ] + assert caught.value.task == "leaf" + assert caught.value.model is RouteInputs + + +def test_a_refusal_raised_with_no_report_behind_it_still_answers(): + """The attributes are there to read whether or not pydantic filled them.""" + refusal = TaskInputValidationError("nothing pydantic said") + assert refusal.errors == [] + assert refusal.task == "" + assert refusal.model is None + + +def test_a_socket_refusal_is_still_a_value_error(): + """The carrier is added to the refusal, not put in place of it.""" + from node_graph.errors import SocketValueError + + assert issubclass(SocketValueError, ValueError) + with pytest.raises(ValueError): + route.build(parameters={"CONTROL": {"calculation": "scf"}}) + + +# -------------------------------------------------------------------------- +# 25. A graph body is handed the types its model declares +# -------------------------------------------------------------------------- + + +#: What the graph bodies below were handed, read back after the build. +BODY_SAW: dict = {} + + +class SpinInputs(BaseModel): + spin: Color = Color.RED + ratio: Decimal = Decimal("0.10") + + +@task.graph(input_model=SpinInputs) +def spun(spin, ratio): + BODY_SAW["spin"] = spin + BODY_SAW["ratio"] = ratio + return add(x=1, y=2) + + +def test_a_graph_body_is_handed_the_member_its_field_declares(): + """The body runs the model's rules, so it should run on the model's values.""" + BODY_SAW.clear() + spun.build(spin="blue", ratio="0.25") + assert BODY_SAW["spin"] == Color.BLUE + assert BODY_SAW["spin"] in (Color.RED, Color.BLUE) + assert BODY_SAW["ratio"] == Decimal("0.25") + + +def test_the_rebuilt_value_still_carries_the_tag_its_link_is_drawn_from(): + """The control: rebuilding a value must not cost the body its link.""" + from node_graph.socket import TaggedValue + + BODY_SAW.clear() + graph = spun.build(spin="blue", ratio="0.25") + assert isinstance(BODY_SAW["spin"], TaggedValue) + assert "add.outputs.result -> graph_outputs.outputs.result" in links_of(graph) + + +class NestedSpin(BaseModel): + spin: Color = Color.RED + + +class HoldsNestedSpin(BaseModel): + inner: NestedSpin = NestedSpin() + ratio: Decimal = Decimal("0.10") + + +@task.graph(input_model=HoldsNestedSpin) +def nested_spun(inner, ratio): + BODY_SAW["inner"] = inner + return add(x=1, y=2) + + +def test_a_member_below_a_namespace_is_rebuilt_and_still_tagged(): + """A namespace's members carry tags of their own, so the walk goes in.""" + from node_graph.socket import TaggedValue + + BODY_SAW.clear() + nested_spun.build(inner={"spin": "blue"}, ratio="0.25") + assert BODY_SAW["inner"]["spin"] == Color.BLUE + assert isinstance(BODY_SAW["inner"]["spin"], TaggedValue) + + +# -------------------------------------------------------------------------- +# 26. A model that names itself +# -------------------------------------------------------------------------- + + +class Chain(BaseModel): + """A model whose field names the model itself.""" + + n: int = 0 + next: Optional["Chain"] = None + + @field_validator("n") + @classmethod + def _small(cls, value): + if value > 10: + raise ValueError("n is capped at 10") + return value + + +Chain.model_rebuild() + + +def test_a_model_that_names_itself_is_refused_where_it_is_declared(): + """A namespace holds one socket per field, and this one has no bottom. + + Nothing else could answer: the walk that builds the sockets, and every + rebuild the checkpoints make, would each follow the chain forever. + """ + with pytest.raises(ModelContractError, match="Chain -> Chain"): + + @task(input_model=Chain) + def chained(n, next): + return n + + +def test_the_same_model_one_link_shorter_is_a_namespace_with_a_bottom(): + """The control: what is refused is the cycle, not the nesting.""" + + class Link(BaseModel): + n: int = 0 + + class Head(BaseModel): + n: int = 0 + next: Optional[Link] = None + + @task(input_model=Head) + def headed(n, next): + return n + + graph = Graph(name="headed") + assert graph.add_task(headed, "h", n=1, next={"n": 2}) is not None + + +def test_the_shadows_a_write_is_judged_by_end_at_the_same_bound(): + """The rebuilds terminate on their own, so no other entry point can hang.""" + validate_wiring_inputs(Chain, {"n": 1, "next": {"n": 2}}, label="c", complete=False) + with pytest.raises(TaskInputValidationError, match="n is capped at 10"): + validate_wiring_inputs(Chain, {"n": 99}, label="c", complete=False) + + +# -------------------------------------------------------------------------- +# 27. A handle that does not take the name of the function it decorates +# -------------------------------------------------------------------------- + + +class BandCount(BaseModel): + """A rule no socket type check can stand in for.""" + + nbnd: int = 1 + + @field_validator("nbnd") + @classmethod + def _counted(cls, value): + if value < 0: + raise ValueError("nbnd counts bands, so it cannot be negative") + return value + + +class BandCap(BaseModel): + """The opposite rule on the same field, for a second task off one function.""" + + nbnd: int = 1 + + @field_validator("nbnd") + @classmethod + def _capped(cls, value): + if value > 10: + raise ValueError("nbnd is capped at 10") + return value + + +def _band_body(nbnd): + return nbnd + + +#: The decorated name is never rebound, so the module still binds the function. +counted_bands = task(input_model=BandCount)(_band_body) +capped_bands = task(input_model=BandCap)(_band_body) + + +@task(input_model=BandCount) +def rebound_bands(nbnd): + """The ordinary spelling: the handle replaces the module global.""" + return nbnd + + +@pytest.mark.parametrize( + "handle", + [counted_bands, rebound_bands], + ids=["handle named apart", "handle rebound to the name"], +) +def test_the_run_edge_holds_the_model_whatever_the_handle_is_called(handle): + """The body runs through the executor, and the executor must reach the wrapper.""" + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + handle.run(nbnd=-5) + + +@pytest.mark.parametrize( + "handle", + [counted_bands, rebound_bands], + ids=["handle named apart", "handle rebound to the name"], +) +def test_the_write_reads_the_model_by_the_same_route(handle): + """The write checkpoint finds the model through the same executor.""" + graph = Graph(name="named_apart") + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + graph.add_task(handle, "b", nbnd=-5) + + +def test_the_function_itself_is_left_unstamped(): + """What is bound is the wrapper: the function still enforces nothing.""" + from node_graph.input_model import input_model_of_callable + + assert input_model_of_callable(_band_body) is None + assert input_model_of_callable(counted_bands._spec.executor.callable) is BandCount + + +def test_the_wrapper_keeps_the_name_a_process_is_labelled_with(): + """Only the name the executor records changes, not the name of the callable.""" + executor = counted_bands._spec.executor + assert executor.callable_name != "_band_body" + assert executor.callable.__name__ == "_band_body" + + +def test_two_tasks_off_one_function_each_keep_their_own_model(): + """Two wrappers cannot answer to one name, so each is bound under its own.""" + assert counted_bands.run(nbnd=20) == 20 + with pytest.raises(TaskInputValidationError, match="capped at 10"): + capped_bands.run(nbnd=20) + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + counted_bands.run(nbnd=-5) + + +def test_a_spec_rebuilt_from_its_stored_form_still_carries_the_model(): + """What is stored is a module path, so the path has to name what enforces it.""" + from node_graph.input_model import _input_model_of_task + from node_graph.task_spec import TaskSpec + + rebuilt = TaskSpec.from_dict(counted_bands._spec.to_dict()).to_task(name="b") + assert _input_model_of_task(rebuilt) is BandCount + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + rebuilt.set_inputs({"nbnd": -5}) + + +def test_the_run_edge_takes_what_an_engine_passes_positionally(): + """An engine that calls the body positionally addresses the same fields. + + AiiDA's process functions call the callable they wrap that way, so a + keyword-only wrapper would fail before the model ever ran. + """ + from node_graph.input_model import validated_callable + + ran = validated_callable(_band_body, input_model=BandCount) + assert ran(20) == 20 + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + ran(-5) + + +def _rebound_body(nbnd): + """Stand for a body an engine's own decorator wraps after validation.""" + return nbnd + + +def test_a_wrapper_over_the_validated_one_takes_over_its_name(): + """An engine's decorator is what the executor resolves to; the checks stay inside.""" + import functools + import sys + + from node_graph.executor import BOUND_NAME_ATTR + from node_graph.input_model import ( + input_model_of_callable, + rebind_executor_callable, + validated_callable, + ) + + inner = validated_callable(_rebound_body, input_model=BandCount) + + @functools.wraps(inner) + def outer(*args, **kwargs): + return inner(*args, **kwargs) + + rebind_executor_callable(outer, inner) + bound = getattr(outer, BOUND_NAME_ATTR) + assert getattr(sys.modules[_rebound_body.__module__], bound) is outer + assert input_model_of_callable(outer) is BandCount + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + outer(nbnd=-5) + + +def test_a_wrapper_stored_by_value_carries_its_model_with_it(): + """The other route out of this module: nothing to bind, and nothing lost.""" + + def nested(nbnd): + return nbnd + + handle = task(input_model=BandCount)(nested) + assert handle._spec.executor.mode.value == "pickled_callable" + with pytest.raises(TaskInputValidationError, match="cannot be negative"): + handle.run(nbnd=-5) + + +# -------------------------------------------------------------------------- +# 28. What a field nobody wrote is worth by the time the body runs +# -------------------------------------------------------------------------- + + +class System(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 Route(BaseModel): + spin: Color = Color.RED + system: System = System() + + +@task(input_model=Route) +def route_body(spin, system): + BODY_SAW["leaf"] = {"spin": spin, "system": dict(system)} + return 1 + + +@task.graph(input_model=Route) +def route_graph(spin, system): + BODY_SAW["graph"] = {"spin": spin, "system": dict(system)} + return add(x=1, y=2) + + +def test_the_body_is_handed_the_members_that_were_written_and_no_others(): + """A member left out is nothing to hand on; a field left out the model answers for.""" + BODY_SAW.clear() + route_body.run(system={"nbnd": 20}) + assert BODY_SAW["leaf"]["system"] == {"nbnd": 20} + assert BODY_SAW["leaf"]["spin"] == Color.RED + + +def test_the_graph_body_reads_the_same_write_the_same_way(): + """The sockets decide what is collected, so both bodies see one write alike.""" + BODY_SAW.clear() + route_graph.build(system={"nbnd": 20}) + assert BODY_SAW["graph"]["system"] == {"nbnd": 20} + assert BODY_SAW["graph"]["spin"] == Color.RED + + +def test_a_namespace_nobody_wrote_at_all_reaches_the_body_empty(): + """A field left out is a namespace with nothing written into it, not a hole.""" + BODY_SAW.clear() + route_body.run(spin=Color.BLUE) + assert BODY_SAW["leaf"]["system"] == {} + assert BODY_SAW["leaf"]["spin"] == Color.BLUE + + +def test_the_graph_body_is_handed_the_same_empty_namespace(): + """The graph checkpoint answers for every field, not only those that arrived.""" + BODY_SAW.clear() + route_graph.build(spin=Color.BLUE) + assert BODY_SAW["graph"]["system"] == {} + + +def test_a_mapping_nobody_wrote_reaches_the_body_empty_too(): + """Same rule, same shape: a mapping left out has no keys to hand over.""" + BODY_SAW.clear() + blocks_body.run() + assert BODY_SAW["blocks"] == {} + + +class Structured(BaseModel): + """A member with no default, which a write cannot leave to the model.""" + + structure: str + nbnd: int = 1 + + +class HoldsStructured(BaseModel): + system: Structured = Structured(structure="si") + + +@task(input_model=HoldsStructured) +def structured_body(system): + return 1 + + +def test_a_member_with_no_default_is_still_refused_when_it_is_absent(): + """Sockets free to be missing are how a write may be partial, not the model.""" + with pytest.raises(TaskInputValidationError, match="Field required"): + structured_body.run(system={"nbnd": 3}) + + +class Coupled(BaseModel): + """A namespace whose rule reads a member the write need not name.""" + + nbnd: int = 1 + nosym: bool = False + + @model_validator(mode="after") + def _large_counts_need_nosym(self): + if self.nbnd > 10 and not self.nosym: + raise ValueError("a large band count needs nosym") + return self + + +class HoldsCoupled(BaseModel): + system: Coupled = Coupled() + + +@task(input_model=HoldsCoupled) +def coupled_body(system): + return 1 + + +def test_a_rule_reading_an_unwritten_member_sees_its_default_and_may_refuse(): + """What the body is handed is narrowed; what the model is given is not.""" + with pytest.raises(TaskInputValidationError, match="needs nosym"): + coupled_body.run(system={"nbnd": 20}) + + +def test_the_same_write_naming_the_member_the_rule_reads_is_accepted(): + """The control: the rule is answered on the write, not skipped by it.""" + assert coupled_body.run(system={"nbnd": 20, "nosym": True}) == 1 + + +class Block(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 Blocks(BaseModel): + blocks: dict[str, Block] = {} + + +@task(input_model=Blocks) +def blocks_body(blocks): + BODY_SAW["blocks"] = {key: dict(item) for key, item in blocks.items()} + return 1 + + +def test_a_mapping_hands_each_item_over_on_the_same_terms(): + """A mapping is one namespace per key, so the rule holds one level down.""" + BODY_SAW.clear() + blocks_body.run(blocks={"occ_1": {"num_iter": 42}}) + assert BODY_SAW["blocks"] == {"occ_1": {"num_iter": 42}} + + +class LabelledBlock(BaseModel): + """An item member with no default, which a write cannot leave out.""" + + label: str + num_iter: int = 100 + + +class LabelledBlocks(BaseModel): + blocks: dict[str, LabelledBlock] = {} + + +@task(input_model=LabelledBlocks) +def labelled_blocks_body(blocks): + return 1 + + +def test_an_item_member_with_no_default_is_still_refused_when_absent(): + """The model reads the whole item, whatever the sockets collect.""" + with pytest.raises(TaskInputValidationError, match="Field required"): + labelled_blocks_body.run(blocks={"occ_1": {"num_iter": 42}}) + + +class CoupledBlock(BaseModel): + """An item whose rule reads a member the write need not name.""" + + num_iter: int = 100 + frozen: bool = False + + @model_validator(mode="after") + def _long_runs_need_freezing(self): + if self.num_iter > 500 and not self.frozen: + raise ValueError("a long run needs frozen") + return self + + +class CoupledBlocks(BaseModel): + blocks: dict[str, CoupledBlock] = {} + + +@task(input_model=CoupledBlocks) +def coupled_blocks_body(blocks): + return 1 + + +def test_an_items_rule_reads_the_member_the_write_left_out(): + """The narrowing is what the body is handed, not what the item is judged as.""" + with pytest.raises(TaskInputValidationError, match="needs frozen"): + coupled_blocks_body.run(blocks={"occ_1": {"num_iter": 900}}) + + +class Rendered(BaseModel): + """A member whose serializer decides its stored form, not its own type.""" + + value: Decimal = Decimal("0.10") + + @field_serializer("value") + def _dump(self, value: Decimal) -> str: + return str(value) + + +class HoldsRendered(BaseModel): + inner: Rendered = Rendered() + + +@task(input_model=HoldsRendered) +def rendered_body(inner): + BODY_SAW["rendered"] = inner + return 1 + + +def test_a_written_member_reaches_the_body_in_the_type_its_model_declares(): + """Narrowing to what was written is not rendering it as storage would.""" + BODY_SAW.clear() + rendered_body.run(inner={"value": "1.50"}) + assert BODY_SAW["rendered"] == {"value": Decimal("1.50")} + + +# -------------------------------------------------------------------------- +# 29. A write that names some of a namespace's members +# -------------------------------------------------------------------------- + + +class Pair(BaseModel): + """Two members with no default, one of which a link may deliver.""" + + n: int + m: int + + +class HoldsPair(BaseModel): + pair: Pair + top: int = 0 + + +@task(input_model=HoldsPair) +def paired(pair, top): + return 1 + + +def test_a_write_may_name_some_of_a_namespaces_members(): + """A member arriving by a later link is as ordinary as a field arriving by one.""" + graph = Graph(name="partial_nested") + assert graph.add_task(paired, "p", pair={"n": 2}) is not None + + +def test_the_same_write_at_the_top_level_is_the_reference(): + """The control: the outer fields have always been free to be written apart.""" + graph = Graph(name="partial_top") + assert graph.add_task(paired, "p", top=3) is not None + + +def test_what_a_partial_write_does_name_is_still_checked(): + """Free to be missing is not free to be anything.""" + graph = Graph(name="partial_typed") + with pytest.raises(TaskInputValidationError) as caught: + graph.add_task(paired, "p", pair={"n": "two"}) + assert [error["loc"] for error in caught.value.errors] == [("pair", "n")] + + +def test_a_call_that_names_everything_it_will_write_still_wants_it_all(): + """The other checkpoint is unchanged: a call is the whole of what is written.""" + with pytest.raises(TaskInputValidationError, match="Field required"): + validate_wiring_inputs(HoldsPair, {"pair": {"n": 2}}, label="p")