Skip to content

Commit

Permalink
Adding a structured object enum fallback
Browse files Browse the repository at this point in the history
  • Loading branch information
blast-hardcheese committed Jan 31, 2025
1 parent d21d8fc commit 84a0d37
Show file tree
Hide file tree
Showing 4 changed files with 257 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,12 @@
NeedsenumInput,
NeedsenumOutput,
)
from .needsEnumObject import (
encode_NeedsenumobjectInput,
NeedsenumobjectOutput,
NeedsenumobjectInput,
NeedsenumobjectErrors,
)


class EnumserviceService:
Expand All @@ -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,
)
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
Literal,
Optional,
Mapping,
NewType,
NotRequired,
Union,
Tuple,
Expand All @@ -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
)
Original file line number Diff line number Diff line change
@@ -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
99 changes: 97 additions & 2 deletions tests/codegen/snapshot/test_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit 84a0d37

Please sign in to comment.