feat: recursive types

Author: darshkpatel

src/replit_river/codegen/client.py

@@ -101,6 +101,8 @@ class RiverConcreteType(BaseModel):
     items: "RiverType | None" = Field(default=None)
     const: str | int | None = Field(default=None)
     patternProperties: dict[str, "RiverType"] = Field(default_factory=lambda: dict())
+    id_: str | None = Field(default=None, alias="$id")
+    ref: str | None = Field(default=None, alias="$ref")
 
 
 class RiverUnionType(BaseModel):
@@ -158,7 +160,12 @@ def encode_type(
     base_model: str,
     in_module: list[ModuleName],
     permit_unknown_members: bool,
+    type_registry: dict[str, TypeName] | None = None,
 ) -> tuple[TypeExpression, list[ModuleName], list[FileContents], set[TypeName]]:
+    # Registry to track $id -> TypeName mappings for resolving $ref
+    if type_registry is None:
+        type_registry = {}
+
     def _make_open_union_type_expr(one_of: list[TypeExpression]) -> OpenUnionTypeExpr:
         if base_model == "RiverError":
             return OpenUnionTypeExpr(
@@ -175,6 +182,17 @@ def _make_open_union_type_expr(one_of: list[TypeExpression]) -> OpenUnionTypeExp
 
     encoder_name: TypeName | None = None  # defining this up here to placate mypy
     chunks: list[FileContents] = []
+
+    # Handle $ref - return a forward reference to the registered type
+    if isinstance(type, RiverConcreteType) and type.ref is not None:
+        ref_id = type.ref
+        if ref_id in type_registry:
+            # Use forward reference string for the type
+            return (TypeName(f'"{type_registry[ref_id].value}"'), [], [], set())
+        else:
+            # Unknown ref, fall back to Any
+            return (TypeName("Any"), [], [], set())
+
     if isinstance(type, RiverNotType):
         return (NoneTypeExpr(), [], [], set())
     elif isinstance(type, RiverUnionType):
@@ -269,6 +287,7 @@ def flatten_union(tpe: RiverType) -> list[RiverType]:
                                 base_model,
                                 in_module,
                                 permit_unknown_members=permit_unknown_members,
+                                type_registry=type_registry,
                             )
                             one_of.append(type_name)
                             chunks.extend(contents)
@@ -304,6 +323,7 @@ def flatten_union(tpe: RiverType) -> list[RiverType]:
                             base_model,
                             in_module,
                             permit_unknown_members=permit_unknown_members,
+                            type_registry=type_registry,
                         )
                         one_of.append(type_name)
                         chunks.extend(contents)
@@ -377,6 +397,7 @@ def {_field_name}(
                 base_model,
                 in_module,
                 permit_unknown_members=permit_unknown_members,
+                type_registry=type_registry,
             )
             any_of.append(type_name)
             chunks.extend(contents)
@@ -462,6 +483,7 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
             base_model,
             in_module,
             permit_unknown_members=permit_unknown_members,
+            type_registry=type_registry,
         )
     elif isinstance(type, RiverConcreteType):
         typeddict_encoder = list[str]()
@@ -509,6 +531,7 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
                 base_model,
                 in_module,
                 permit_unknown_members=permit_unknown_members,
+                type_registry=type_registry,
             )
             typeddict_encoder.append("TODO: dstewart")
             return (ListTypeExpr(type_name), module_info, type_chunks, encoder_names)
@@ -523,10 +546,15 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
                 base_model,
                 in_module,
                 permit_unknown_members=permit_unknown_members,
+                type_registry=type_registry,
             )
             return (DictTypeExpr(type_name), module_info, type_chunks, encoder_names)
         assert type.type == "object", type.type
 
+        # Register $id for this type so $ref can resolve to it
+        if type.id_ is not None:
+            type_registry[type.id_] = prefix
+
         current_chunks: list[str] = [
             f"class {render_literal_type(prefix)}({base_model}):"
         ]
@@ -551,6 +579,7 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
                     "BaseModel" if base_model == "RiverError" else base_model,
                     in_module,
                     permit_unknown_members=permit_unknown_members,
+                    type_registry=type_registry,
                 )
                 encoder_name = None
                 chunks.extend(contents)

tests/v1/codegen/snapshot/snapshots/test_recursive_types/__init__.py

@@ -0,0 +1,13 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from pydantic import BaseModel
+from typing import Literal
+
+import replit_river as river
+
+
+from .recursiveService import RecursiveserviceService
+
+
+class RecursiveClient:
+    def __init__(self, client: river.Client[Literal[None]]):
+        self.recursiveService = RecursiveserviceService(client)

tests/v1/codegen/snapshot/snapshots/test_recursive_types/recursiveService/__init__.py

@@ -0,0 +1,41 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+from typing import Any
+import datetime
+
+from pydantic import TypeAdapter
+
+from replit_river.error_schema import RiverError, RiverErrorTypeAdapter
+import replit_river as river
+
+
+from .getTree import (
+    GettreeInput,
+    GettreeOutput,
+    GettreeOutputTypeAdapter,
+    encode_GettreeInput,
+)
+
+
+class RecursiveserviceService:
+    def __init__(self, client: river.Client[Any]):
+        self.client = client
+
+    async def getTree(
+        self,
+        input: GettreeInput,
+        timeout: datetime.timedelta,
+    ) -> GettreeOutput:
+        return await self.client.send_rpc(
+            "recursiveService",
+            "getTree",
+            input,
+            encode_GettreeInput,
+            lambda x: GettreeOutputTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            lambda x: RiverErrorTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            timeout,
+        )

tests/v1/codegen/snapshot/snapshots/test_recursive_types/recursiveService/getTree.py

@@ -0,0 +1,49 @@
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+import datetime
+from typing import (
+    Any,
+    Literal,
+    Mapping,
+    NotRequired,
+    TypedDict,
+)
+from typing_extensions import Annotated
+
+from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
+from replit_river.error_schema import RiverError
+from replit_river.client import (
+    RiverUnknownError,
+    translate_unknown_error,
+    RiverUnknownValue,
+    translate_unknown_value,
+)
+
+import replit_river as river
+
+
+def encode_GettreeInput(
+    x: "GettreeInput",
+) -> Any:
+    return {
+        k: v
+        for (k, v) in (
+            {
+                "rootId": x.get("rootId"),
+            }
+        ).items()
+        if v is not None
+    }
+
+
+class GettreeInput(TypedDict):
+    rootId: str
+
+
+class GettreeOutput(BaseModel):
+    children: list["GettreeOutput"]
+    id: str
+    name: str
+
+
+GettreeOutputTypeAdapter: TypeAdapter[GettreeOutput] = TypeAdapter(GettreeOutput)

tests/v1/codegen/snapshot/test_recursive.py

@@ -0,0 +1,15 @@
+from pytest_snapshot.plugin import Snapshot
+
+from tests.fixtures.codegen_snapshot_fixtures import validate_codegen
+
+
+async def test_recursive_types(snapshot: Snapshot) -> None:
+    """Test that recursive types using $id/$ref are generated correctly."""
+    validate_codegen(
+        snapshot=snapshot,
+        snapshot_dir="tests/v1/codegen/snapshot/snapshots",
+        read_schema=lambda: open("tests/v1/codegen/types/recursive_schema.json"),
+        target_path="test_recursive_types",
+        client_name="RecursiveClient",
+        protocol_version="v1.1",
+    )

tests/v1/codegen/types/recursive_schema.json

@@ -0,0 +1,42 @@
+{
+  "services": {
+    "recursiveService": {
+      "procedures": {
+        "getTree": {
+          "type": "rpc",
+          "input": {
+            "type": "object",
+            "properties": {
+              "rootId": {
+                "type": "string"
+              }
+            },
+            "required": ["rootId"]
+          },
+          "output": {
+            "$id": "TreeNode",
+            "type": "object",
+            "properties": {
+              "id": {
+                "type": "string"
+              },
+              "name": {
+                "type": "string"
+              },
+              "children": {
+                "type": "array",
+                "items": {
+                  "$ref": "TreeNode"
+                }
+              }
+            },
+            "required": ["id", "name", "children"]
+          },
+          "errors": {
+            "not": {}
+          }
+        }
+      }
+    }
+  }
+}