Adding a structured object enum fallback

Author: blast-hardcheese

tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/__init__.py

@@ -15,6 +15,12 @@
     NeedsenumInput,
     NeedsenumOutput,
 )
+from .needsEnumObject import (
+    encode_NeedsenumobjectInput,
+    NeedsenumobjectOutput,
+    NeedsenumobjectInput,
+    NeedsenumobjectErrors,
+)
 
 
 class EnumserviceService:
@@ -39,3 +45,22 @@ async def needsEnum(
             ),
             timeout,
         )
+
+    async def needsEnumObject(
+        self,
+        input: NeedsenumobjectInput,
+        timeout: datetime.timedelta,
+    ) -> NeedsenumobjectOutput:
+        return await self.client.send_rpc(
+            "enumService",
+            "needsEnumObject",
+            input,
+            encode_NeedsenumobjectInput,
+            lambda x: TypeAdapter(NeedsenumobjectOutput).validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            lambda x: TypeAdapter(NeedsenumobjectErrors).validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            timeout,
+        )

tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py

@@ -10,6 +10,7 @@
     Literal,
     Optional,
     Mapping,
+    NewType,
     NotRequired,
     Union,
     Tuple,
@@ -24,5 +25,11 @@
 
 NeedsenumInput = Literal["in_first"] | Literal["in_second"]
 encode_NeedsenumInput: Callable[["NeedsenumInput"], Any] = lambda x: x
-NeedsenumOutput = Literal["out_first"] | Literal["out_second"]
-NeedsenumErrors = Literal["err_first"] | Literal["err_second"]
+NeedsenumOutputAnyOf__Unknown = NewType("NeedsenumOutputAnyOf__Unknown", object)
+NeedsenumOutput = (
+    Literal["out_first"] | Literal["out_second"] | NeedsenumOutputAnyOf__Unknown
+)
+NeedsenumErrorsAnyOf__Unknown = NewType("NeedsenumErrorsAnyOf__Unknown", object)
+NeedsenumErrors = (
+    Literal["err_first"] | Literal["err_second"] | NeedsenumErrorsAnyOf__Unknown
+)

tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py

@@ -0,0 +1,126 @@
+# ruff: noqa
+# Code generated by river.codegen. DO NOT EDIT.
+from collections.abc import AsyncIterable, AsyncIterator
+import datetime
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Mapping,
+    NewType,
+    NotRequired,
+    Union,
+    Tuple,
+    TypedDict,
+)
+
+from pydantic import BaseModel, Field, TypeAdapter
+from replit_river.error_schema import RiverError
+
+import replit_river as river
+
+
+encode_NeedsenumobjectInputOneOf_in_first: Callable[
+    ["NeedsenumobjectInputOneOf_in_first"], Any
+] = lambda x: {
+    k: v
+    for (k, v) in (
+        {
+            "$kind": x.get("kind"),
+            "value": x.get("value"),
+        }
+    ).items()
+    if v is not None
+}
+
+
+class NeedsenumobjectInputOneOf_in_first(TypedDict):
+    kind: Literal["in_first"]
+    value: str
+
+
+encode_NeedsenumobjectInputOneOf_in_second: Callable[
+    ["NeedsenumobjectInputOneOf_in_second"], Any
+] = lambda x: {
+    k: v
+    for (k, v) in (
+        {
+            "$kind": x.get("kind"),
+            "bleep": x.get("bleep"),
+        }
+    ).items()
+    if v is not None
+}
+
+
+class NeedsenumobjectInputOneOf_in_second(TypedDict):
+    kind: Literal["in_second"]
+    bleep: int
+
+
+NeedsenumobjectInput = (
+    NeedsenumobjectInputOneOf_in_first | NeedsenumobjectInputOneOf_in_second
+)
+
+encode_NeedsenumobjectInput: Callable[["NeedsenumobjectInput"], Any] = (
+    lambda x: encode_NeedsenumobjectInputOneOf_in_first(x)
+    if x["kind"] == "in_first"
+    else encode_NeedsenumobjectInputOneOf_in_second(x)
+)
+
+
+class NeedsenumobjectOutputFooOneOf_out_first(BaseModel):
+    kind: Literal["out_first"] = Field(
+        "out_first",
+        alias="$kind",  # type: ignore
+    )
+
+    foo: int
+
+
+class NeedsenumobjectOutputFooOneOf_out_second(BaseModel):
+    kind: Literal["out_second"] = Field(
+        "out_second",
+        alias="$kind",  # type: ignore
+    )
+
+    bar: int
+
+
+NeedsenumobjectOutputFooAnyOf__Unknown = NewType(
+    "NeedsenumobjectOutputFooAnyOf__Unknown", object
+)
+NeedsenumobjectOutputFoo = (
+    NeedsenumobjectOutputFooOneOf_out_first
+    | NeedsenumobjectOutputFooOneOf_out_second
+    | NeedsenumobjectOutputFooAnyOf__Unknown
+)
+
+
+class NeedsenumobjectOutput(BaseModel):
+    foo: Optional[NeedsenumobjectOutputFoo] = None
+
+
+class NeedsenumobjectErrorsFooAnyOf_0(RiverError):
+    beep: Optional[Literal["err_first"]] = None
+
+
+class NeedsenumobjectErrorsFooAnyOf_1(RiverError):
+    borp: Optional[Literal["err_second"]] = None
+
+
+NeedsenumobjectErrorsFooAnyOf__Unknown = NewType(
+    "NeedsenumobjectErrorsFooAnyOf__Unknown", object
+)
+NeedsenumobjectErrorsFoo = (
+    NeedsenumobjectErrorsFooAnyOf_0
+    | NeedsenumobjectErrorsFooAnyOf_1
+    | NeedsenumobjectErrorsFooAnyOf__Unknown
+)
+
+
+class NeedsenumobjectErrors(RiverError):
+    foo: Optional[NeedsenumobjectErrorsFoo] = None

tests/codegen/snapshot/test_enum.py

@@ -6,7 +6,7 @@
 
 from replit_river.codegen.client import schema_to_river_client_codegen
 
-test_unknown_enum_schema: str = """
+test_unknown_enum_schema = """
 {
   "services": {
     "enumService": {
@@ -49,12 +49,107 @@
               }
             ]
           }
+        },
+        "needsEnumObject": {
+          "type": "rpc",
+          "input": {
+            "anyOf": [
+              {
+                "type": "object",
+                "properties": {
+                  "$kind": {
+                    "const": "in_first",
+                    "type": "string"
+                  },
+                  "value": {
+                    "type": "string"
+                  }
+                },
+                "required": ["$kind", "value"]
+              },
+              {
+                "type": "object",
+                "properties": {
+                  "$kind": {
+                    "const": "in_second",
+                    "type": "string"
+                  },
+                  "bleep": {
+                    "type": "integer"
+                  }
+                },
+                "required": ["$kind", "bleep"]
+              }
+            ]
+          },
+          "output": {
+            "type": "object",
+            "properties": {
+              "foo": {
+                "anyOf": [
+                  {
+                    "type": "object",
+                    "properties": {
+                      "$kind": {
+                        "const": "out_first",
+                        "type": "string"
+                      },
+                      "foo": {
+                        "type": "integer"
+                      }
+                    },
+                    "required": ["$kind", "foo"]
+                  },
+                  {
+                    "type": "object",
+                    "properties": {
+                      "$kind": {
+                        "const": "out_second",
+                        "type": "string"
+                      },
+                      "bar": {
+                        "type": "integer"
+                      }
+                    },
+                    "required": ["$kind", "bar"]
+                  }
+                ]
+              }
+            }
+          },
+          "errors": {
+            "type": "object",
+            "properties": {
+              "foo": {
+                "anyOf": [
+                  {
+                    "type": "object",
+                    "properties": {
+                      "beep": {
+                        "type": "string",
+                        "const": "err_first"
+                      }
+                    }
+                  },
+                  {
+                    "type": "object",
+                    "properties": {
+                      "borp": {
+                        "type": "string",
+                        "const": "err_second"
+                      }
+                    }
+                  }
+                ]
+              }
+            }
+          }
         }
       }
     }
   }
 }
-    """
+"""
 
 
 class UnclosableStringIO(StringIO):