✨ Take a task's input and output contracts from Pydantic models

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

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

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)