feat: generalize draft schema projection

src/wf_api/schema_projection.py

@@ -8,6 +8,51 @@ from jsonschema import Draft202012Validator, SchemaError
 JsonObject = dict[str, Any]
 
 
+def project_property_to_schema_path(
+    *,
+    target_schema: JsonObject,
+    source_schema: JsonObject,
+    source_field: str,
+    target_parts: tuple[str, ...],
+) -> JsonObject:
+    """Copy one source property schema into a target JSON Schema object path."""
+    if not target_parts:
+        raise ValueError("target schema path must not be empty")
+    _check_schema("target_schema", target_schema)
+    _check_schema("source_schema", source_schema)
+    source_properties = source_schema.get("properties")
+    if not isinstance(source_properties, dict) or source_field not in source_properties:
+        raise ValueError(f"source field {source_field!r} is not declared")
+    source_property = source_properties[source_field]
+    if not isinstance(source_property, dict):
+        raise ValueError(f"source field {source_field!r} is not a JSON Schema object")
+
+    projected = deepcopy(target_schema)
+    _ensure_object_schema(projected, "target_schema")
+    parent = projected
+    for index, part in enumerate(target_parts[:-1]):
+        properties = _properties_for_object(parent, ".".join(target_parts[:index]) or "target_schema")
+        child = properties.get(part)
+        if child is None:
+            child = {"type": "object", "properties": {}}
+            properties[part] = child
+        if not isinstance(child, dict):
+            raise ValueError(f"schema path {'.'.join(target_parts[: index + 1])!r} is not an object")
+        _ensure_object_schema(child, ".".join(target_parts[: index + 1]))
+        parent = child
+
+    properties = _properties_for_object(parent, ".".join(target_parts[:-1]) or "target_schema")
+    leaf = target_parts[-1]
+    if leaf in properties:
+        raise ValueError(f"schema path {'.'.join(target_parts)!r} already exists")
+    properties[leaf] = deepcopy(source_property)
+
+    _merge_definition_block(projected, source_schema, "$defs")
+    _merge_definition_block(projected, source_schema, "definitions")
+    _check_schema("projected target_schema", projected)
+    return projected
+
+
 def project_output_property_to_state_schema(
     *,
     state_schema: JsonObject,
@@ -15,38 +60,32 @@ def project_output_property_to_state_schema(
     output_field: str,
     state_field: str,
 ) -> JsonObject:
-    """Project one capability output property schema into workflow state schema.
+    """Root state projection convenience wrapper.
 
-    Capability output schemas may use local references such as
+    Preserves the original error message wording for backward compatibility.
-    ``{"$ref": "#/$defs/Snapshot"}``. Copying only the property schema would
-    create dangling references, so this helper also merges local definition
-    blocks and rejects conflicting definition names.
     """
-    _check_schema("state_schema", state_schema)
+    try:
-    _check_schema("output_schema", output_schema)
+        return project_property_to_schema_path(
-    output_properties = output_schema.get("properties")
+            target_schema=state_schema,
-    if not isinstance(output_properties, dict) or output_field not in output_properties:
+            source_schema=output_schema,
-        raise ValueError(f"output field {output_field!r} is not declared")
+            source_field=output_field,
-    output_property = output_properties[output_field]
+            target_parts=(state_field,),
-    if not isinstance(output_property, dict):
+        )
-        raise ValueError(f"output field {output_field!r} is not a JSON Schema object")
+    except ValueError as exc:
-
+        msg = str(exc)
-    projected = deepcopy(state_schema)
+        if msg.startswith("source field ") and "is not declared" in msg:
-    state_type = projected.get("type")
+            raise ValueError(f"output field {output_field!r} is not declared") from exc
-    if state_type is not None and state_type != "object":
+        if msg.startswith("source field ") and "not a JSON Schema" in msg:
-        raise ValueError("state_schema must be an object schema")
+            raise ValueError(f"output field {output_field!r} is not a JSON Schema object") from exc
-    projected.setdefault("type", "object")
+        if "schema path 'target_schema'" in msg and "is not an object" in msg:
-    properties = projected.setdefault("properties", {})
+            raise ValueError("state_schema must be an object schema") from exc
-    if not isinstance(properties, dict):
+        if msg.startswith("schema path ") and "already exists" in msg:
-        raise ValueError("state_schema.properties must be an object")
+            raise ValueError(f"state field {state_field!r} already exists") from exc
-    if state_field in properties:
+        if "target_schema is not valid JSON Schema" in msg:
-        raise ValueError(f"state field {state_field!r} already exists")
+            raise ValueError(f"state_schema is not valid JSON Schema: {msg.split(': ', 1)[1]}") from exc
-    properties[state_field] = deepcopy(output_property)
+        if "source_schema is not valid JSON Schema" in msg:
-
+            raise ValueError(f"output_schema is not valid JSON Schema: {msg.split(': ', 1)[1]}") from exc
-    _merge_definition_block(projected, output_schema, "$defs")
+        raise
-    _merge_definition_block(projected, output_schema, "definitions")
-    _check_schema("projected state_schema", projected)
-    return projected
 
 
 def _check_schema(name: str, schema: JsonObject) -> None:
@@ -56,6 +95,20 @@ def _check_schema(name: str, schema: JsonObject) -> None:
         raise ValueError(f"{name} is not valid JSON Schema: {exc.message}") from exc
 
 
+def _ensure_object_schema(schema: JsonObject, label: str) -> None:
+    schema_type = schema.get("type")
+    if schema_type is not None and schema_type != "object":
+        raise ValueError(f"schema path {label!r} is not an object")
+    schema.setdefault("type", "object")
+
+
+def _properties_for_object(schema: JsonObject, label: str) -> JsonObject:
+    properties = schema.setdefault("properties", {})
+    if not isinstance(properties, dict):
+        raise ValueError(f"{label}.properties must be an object")
+    return properties
+
+
 def _merge_definition_block(
     target_schema: JsonObject,
     source_schema: JsonObject,

tests/wf_api/test_schema_projection.py

@@ -2,7 +2,10 @@ from __future__ import annotations
 
 import pytest
 
-from wf_api.schema_projection import project_output_property_to_state_schema
+from wf_api.schema_projection import (
+    project_output_property_to_state_schema,
+    project_property_to_schema_path,
+)
 
 
 def test_project_output_property_copies_schema_and_defs() -> None:
@@ -107,3 +110,66 @@ def test_project_output_property_rejects_invalid_output_schema() -> None:
             output_field="after",
             state_field="after",
         )
+
+
+def test_project_schema_property_inserts_nested_path_and_defs() -> None:
+    projected = project_property_to_schema_path(
+        target_schema={"type": "object", "properties": {}},
+        source_schema={
+            "type": "object",
+            "properties": {"after": {"$ref": "#/$defs/Snapshot"}},
+            "$defs": {
+                "Snapshot": {
+                    "type": "object",
+                    "properties": {"clicked": {"type": "boolean"}},
+                }
+            },
+        },
+        source_field="after",
+        target_parts=("session", "after"),
+    )
+
+    assert projected["properties"]["session"]["type"] == "object"
+    assert projected["properties"]["session"]["properties"]["after"] == {
+        "$ref": "#/$defs/Snapshot"
+    }
+    assert projected["$defs"]["Snapshot"]["properties"]["clicked"] == {
+        "type": "boolean"
+    }
+
+
+def test_project_schema_property_rejects_existing_nested_target() -> None:
+    with pytest.raises(ValueError, match="schema path 'session.after' already exists"):
+        project_property_to_schema_path(
+            target_schema={
+                "type": "object",
+                "properties": {
+                    "session": {
+                        "type": "object",
+                        "properties": {"after": {"type": "string"}},
+                    }
+                },
+            },
+            source_schema={
+                "type": "object",
+                "properties": {"after": {"type": "object"}},
+            },
+            source_field="after",
+            target_parts=("session", "after"),
+        )
+
+
+def test_project_schema_property_rejects_non_object_ancestor() -> None:
+    with pytest.raises(ValueError, match="schema path 'session' is not an object"):
+        project_property_to_schema_path(
+            target_schema={
+                "type": "object",
+                "properties": {"session": {"type": "string"}},
+            },
+            source_schema={
+                "type": "object",
+                "properties": {"after": {"type": "object"}},
+            },
+            source_field="after",
+            target_parts=("session", "after"),
+        )