From 20d1287b5891facbb853ac455225bdbb00633c2a Mon Sep 17 00:00:00 2001 From: Devon Stewart Date: Wed, 5 Feb 2025 11:51:23 -0800 Subject: [PATCH 1/6] Better way to model open unions --- src/replit_river/codegen/client.py | 32 ++++++++++-------------------- src/replit_river/codegen/typing.py | 27 ++++++++++++++++--------- 2 files changed, 28 insertions(+), 31 deletions(-) diff --git a/src/replit_river/codegen/client.py b/src/replit_river/codegen/client.py index c58f745c..869b22f4 100644 --- a/src/replit_river/codegen/client.py +++ b/src/replit_river/codegen/client.py @@ -30,11 +30,11 @@ ListTypeExpr, LiteralTypeExpr, ModuleName, + OpenUnionTypeExpr, RenderedPath, TypeExpression, TypeName, UnionTypeExpr, - UnknownTypeExpr, ensure_literal_type, extract_inner_type, render_type_expr, @@ -83,7 +83,6 @@ Literal, Optional, Mapping, - NewType, NotRequired, Union, Tuple, @@ -311,19 +310,12 @@ def flatten_union(tpe: RiverType) -> list[RiverType]: else """, ) + union: TypeExpression if permit_unknown_members: - unknown_name = TypeName(f"{prefix}AnyOf__Unknown") - chunks.append( - FileContents( - f"{unknown_name} = NewType({repr(unknown_name)}, object)" - ) - ) - one_of.append(UnknownTypeExpr(unknown_name)) - chunks.append( - FileContents( - f"{prefix} = {render_type_expr(UnionTypeExpr(one_of))}" - ) - ) + union = OpenUnionTypeExpr(UnionTypeExpr(one_of)) + else: + union = UnionTypeExpr(one_of) + chunks.append(FileContents(f"{prefix} = {render_type_expr(union)}")) chunks.append(FileContents("")) if base_model == "TypedDict": @@ -386,16 +378,12 @@ def flatten_union(tpe: RiverType) -> list[RiverType]: f"encode_{ensure_literal_type(other)}(x)" ) if permit_unknown_members: - unknown_name = TypeName(f"{prefix}AnyOf__Unknown") - chunks.append( - FileContents(f"{unknown_name} = NewType({repr(unknown_name)}, object)") - ) - any_of.append(UnknownTypeExpr(unknown_name)) + union = OpenUnionTypeExpr(UnionTypeExpr(any_of)) + else: + union = UnionTypeExpr(any_of) if is_literal(type): typeddict_encoder = ["x"] - chunks.append( - FileContents(f"{prefix} = {render_type_expr(UnionTypeExpr(any_of))}") - ) + chunks.append(FileContents(f"{prefix} = {render_type_expr(union)}")) if base_model == "TypedDict": encoder_name = TypeName(f"encode_{prefix}") encoder_names.add(encoder_name) diff --git a/src/replit_river/codegen/typing.py b/src/replit_river/codegen/typing.py index af9ae7fa..06513027 100644 --- a/src/replit_river/codegen/typing.py +++ b/src/replit_river/codegen/typing.py @@ -31,8 +31,8 @@ class UnionTypeExpr: @dataclass -class UnknownTypeExpr: - name: TypeName +class OpenUnionTypeExpr: + union: UnionTypeExpr TypeExpression = ( @@ -41,7 +41,7 @@ class UnknownTypeExpr: | ListTypeExpr | LiteralTypeExpr | UnionTypeExpr - | UnknownTypeExpr + | OpenUnionTypeExpr ) @@ -55,10 +55,15 @@ def render_type_expr(value: TypeExpression) -> str: return f"Literal[{repr(inner)}]" case UnionTypeExpr(inner): return " | ".join(render_type_expr(x) for x in inner) + case OpenUnionTypeExpr(inner): + return ( + "Annotated[" + f"{render_type_expr(inner)} | RiverUnknownValue," + "WrapValidator(raise_unknown)" + "]" + ) case str(name): return TypeName(name) - case UnknownTypeExpr(name): - return TypeName(name) case other: assert_never(other) @@ -75,10 +80,12 @@ def extract_inner_type(value: TypeExpression) -> TypeName: raise ValueError( f"Attempting to extract from a union, currently not possible: {value}" ) + case OpenUnionTypeExpr(_): + raise ValueError( + f"Attempting to extract from a union, currently not possible: {value}" + ) case str(name): return TypeName(name) - case UnknownTypeExpr(name): - return name case other: assert_never(other) @@ -101,9 +108,11 @@ def ensure_literal_type(value: TypeExpression) -> TypeName: raise ValueError( f"Unexpected expression when expecting a type name: {value}" ) + case OpenUnionTypeExpr(_): + raise ValueError( + f"Unexpected expression when expecting a type name: {value}" + ) case str(name): return TypeName(name) - case UnknownTypeExpr(name): - return name case other: assert_never(other) From deb460aef201c53a70865e757fb1ce46983b7528 Mon Sep 17 00:00:00 2001 From: Devon Stewart Date: Wed, 5 Feb 2025 11:51:38 -0800 Subject: [PATCH 2/6] Fix --output argument description --- src/replit_river/codegen/run.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/replit_river/codegen/run.py b/src/replit_river/codegen/run.py index dd041483..c9ab6384 100644 --- a/src/replit_river/codegen/run.py +++ b/src/replit_river/codegen/run.py @@ -30,7 +30,7 @@ def main() -> None: client = subparsers.add_parser( "client", help="Codegen a River client from JSON schema" ) - client.add_argument("--output", help="output file", required=True) + client.add_argument("--output", help="output path", required=True) client.add_argument("--client-name", help="name of the class", required=True) client.add_argument( "--typed-dict-inputs", From 32aecbb99f3d7f3f5c7c9631169e85fe31c606f1 Mon Sep 17 00:00:00 2001 From: Devon Stewart Date: Wed, 5 Feb 2025 11:53:15 -0800 Subject: [PATCH 3/6] Publish new type for client consumption --- src/replit_river/client.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/replit_river/client.py b/src/replit_river/client.py index c78fb171..d1863e30 100644 --- a/src/replit_river/client.py +++ b/src/replit_river/client.py @@ -7,6 +7,10 @@ from opentelemetry import trace from opentelemetry.trace import Span, SpanKind, Status, StatusCode +from pydantic import ( + BaseModel, + ValidationInfo, +) from replit_river.client_transport import ClientTransport from replit_river.error_schema import RiverError, RiverException @@ -27,6 +31,21 @@ tracer = trace.get_tracer(__name__) +@dataclass(frozen=True) +class RiverUnknownValue(BaseModel): + tag: Literal["RiverUnknownValue"] + value: Any + + +def raise_unknown( + value: Any, handler: Callable[[Any], Any], info: ValidationInfo +) -> Any | RiverUnknownValue: + try: + return handler(value) + except Exception: + return RiverUnknownValue(tag="RiverUnknownValue", value=value) + + class Client(Generic[HandshakeMetadataType]): def __init__( self, From 290e470f90a191a0c33d6be693f739b66b584abd Mon Sep 17 00:00:00 2001 From: Devon Stewart Date: Wed, 5 Feb 2025 16:10:34 -0800 Subject: [PATCH 4/6] Importing newly required stuff --- src/replit_river/codegen/client.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/replit_river/codegen/client.py b/src/replit_river/codegen/client.py index 869b22f4..a56cedd2 100644 --- a/src/replit_river/codegen/client.py +++ b/src/replit_river/codegen/client.py @@ -88,9 +88,11 @@ Tuple, TypedDict, ) +from typing_extensions import Annotated -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError +from replit_river.client import RiverUnknownValue, raise_unknown import replit_river as river From 7275f2e173505e71c6589d975e4b226bf0fd0313 Mon Sep 17 00:00:00 2001 From: Devon Stewart Date: Wed, 5 Feb 2025 16:25:57 -0800 Subject: [PATCH 5/6] Updating tests --- .../rpc/generated/test_service/__init__.py | 2 +- .../rpc/generated/test_service/rpc_method.py | 5 ++-- .../enumService/needsEnum.py | 21 ++++++++-------- .../enumService/needsEnumObject.py | 25 ++++++++----------- .../stream/generated/test_service/__init__.py | 4 +-- .../generated/test_service/stream_method.py | 5 ++-- 6 files changed, 31 insertions(+), 31 deletions(-) diff --git a/tests/codegen/rpc/generated/test_service/__init__.py b/tests/codegen/rpc/generated/test_service/__init__.py index 8012ee6b..fc994615 100644 --- a/tests/codegen/rpc/generated/test_service/__init__.py +++ b/tests/codegen/rpc/generated/test_service/__init__.py @@ -9,7 +9,7 @@ import replit_river as river -from .rpc_method import encode_Rpc_MethodInput, Rpc_MethodInput, Rpc_MethodOutput +from .rpc_method import Rpc_MethodInput, Rpc_MethodOutput, encode_Rpc_MethodInput class Test_ServiceService: diff --git a/tests/codegen/rpc/generated/test_service/rpc_method.py b/tests/codegen/rpc/generated/test_service/rpc_method.py index d6dc64f4..2ec649d3 100644 --- a/tests/codegen/rpc/generated/test_service/rpc_method.py +++ b/tests/codegen/rpc/generated/test_service/rpc_method.py @@ -10,15 +10,16 @@ Literal, Optional, Mapping, - NewType, NotRequired, Union, Tuple, TypedDict, ) +from typing_extensions import Annotated -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError +from replit_river.client import RiverUnknownValue, raise_unknown import replit_river as river diff --git a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py index 4633dfd6..bc1bc434 100644 --- a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py +++ b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py @@ -10,26 +10,27 @@ Literal, Optional, Mapping, - NewType, NotRequired, Union, Tuple, TypedDict, ) +from typing_extensions import Annotated -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError +from replit_river.client import RiverUnknownValue, raise_unknown import replit_river as river NeedsenumInput = Literal["in_first"] | Literal["in_second"] encode_NeedsenumInput: Callable[["NeedsenumInput"], Any] = lambda x: x -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 -) +NeedsenumOutput = Annotated[ + Literal["out_first"] | Literal["out_second"] | RiverUnknownValue, + WrapValidator(raise_unknown), +] +NeedsenumErrors = Annotated[ + Literal["err_first"] | Literal["err_second"] | RiverUnknownValue, + WrapValidator(raise_unknown), +] diff --git a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py index 2fe76987..62c3d484 100644 --- a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py +++ b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py @@ -10,15 +10,16 @@ Literal, Optional, Mapping, - NewType, NotRequired, Union, Tuple, TypedDict, ) +from typing_extensions import Annotated -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError +from replit_river.client import RiverUnknownValue, raise_unknown import replit_river as river @@ -90,14 +91,12 @@ class NeedsenumobjectOutputFooOneOf_out_second(BaseModel): bar: int -NeedsenumobjectOutputFooAnyOf__Unknown = NewType( - "NeedsenumobjectOutputFooAnyOf__Unknown", object -) -NeedsenumobjectOutputFoo = ( +NeedsenumobjectOutputFoo = Annotated[ NeedsenumobjectOutputFooOneOf_out_first | NeedsenumobjectOutputFooOneOf_out_second - | NeedsenumobjectOutputFooAnyOf__Unknown -) + | RiverUnknownValue, + WrapValidator(raise_unknown), +] class NeedsenumobjectOutput(BaseModel): @@ -112,14 +111,12 @@ class NeedsenumobjectErrorsFooAnyOf_1(RiverError): borp: Optional[Literal["err_second"]] = None -NeedsenumobjectErrorsFooAnyOf__Unknown = NewType( - "NeedsenumobjectErrorsFooAnyOf__Unknown", object -) -NeedsenumobjectErrorsFoo = ( +NeedsenumobjectErrorsFoo = Annotated[ NeedsenumobjectErrorsFooAnyOf_0 | NeedsenumobjectErrorsFooAnyOf_1 - | NeedsenumobjectErrorsFooAnyOf__Unknown -) + | RiverUnknownValue, + WrapValidator(raise_unknown), +] class NeedsenumobjectErrors(RiverError): diff --git a/tests/codegen/stream/generated/test_service/__init__.py b/tests/codegen/stream/generated/test_service/__init__.py index 00f9c553..f59c046e 100644 --- a/tests/codegen/stream/generated/test_service/__init__.py +++ b/tests/codegen/stream/generated/test_service/__init__.py @@ -10,9 +10,9 @@ from .stream_method import ( - encode_Stream_MethodInput, - Stream_MethodOutput, Stream_MethodInput, + Stream_MethodOutput, + encode_Stream_MethodInput, ) diff --git a/tests/codegen/stream/generated/test_service/stream_method.py b/tests/codegen/stream/generated/test_service/stream_method.py index d66aff55..75bc2817 100644 --- a/tests/codegen/stream/generated/test_service/stream_method.py +++ b/tests/codegen/stream/generated/test_service/stream_method.py @@ -10,15 +10,16 @@ Literal, Optional, Mapping, - NewType, NotRequired, Union, Tuple, TypedDict, ) +from typing_extensions import Annotated -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError +from replit_river.client import RiverUnknownValue, raise_unknown import replit_river as river From 087e6d116f9a7d4c6da106e8b2f0a24c25a303f8 Mon Sep 17 00:00:00 2001 From: Devon Stewart Date: Thu, 6 Feb 2025 11:08:24 -0800 Subject: [PATCH 6/6] Clarify confusing function name --- src/replit_river/client.py | 2 +- src/replit_river/codegen/client.py | 2 +- src/replit_river/codegen/typing.py | 2 +- tests/codegen/rpc/generated/test_service/rpc_method.py | 2 +- .../snapshots/test_unknown_enum/enumService/needsEnum.py | 6 +++--- .../test_unknown_enum/enumService/needsEnumObject.py | 6 +++--- .../codegen/stream/generated/test_service/stream_method.py | 2 +- 7 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/replit_river/client.py b/src/replit_river/client.py index d1863e30..51e45a2d 100644 --- a/src/replit_river/client.py +++ b/src/replit_river/client.py @@ -37,7 +37,7 @@ class RiverUnknownValue(BaseModel): value: Any -def raise_unknown( +def translate_unknown_value( value: Any, handler: Callable[[Any], Any], info: ValidationInfo ) -> Any | RiverUnknownValue: try: diff --git a/src/replit_river/codegen/client.py b/src/replit_river/codegen/client.py index a56cedd2..f7271dc7 100644 --- a/src/replit_river/codegen/client.py +++ b/src/replit_river/codegen/client.py @@ -92,7 +92,7 @@ from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError -from replit_river.client import RiverUnknownValue, raise_unknown +from replit_river.client import RiverUnknownValue, translate_unknown_value import replit_river as river diff --git a/src/replit_river/codegen/typing.py b/src/replit_river/codegen/typing.py index 06513027..87c577bf 100644 --- a/src/replit_river/codegen/typing.py +++ b/src/replit_river/codegen/typing.py @@ -59,7 +59,7 @@ def render_type_expr(value: TypeExpression) -> str: return ( "Annotated[" f"{render_type_expr(inner)} | RiverUnknownValue," - "WrapValidator(raise_unknown)" + "WrapValidator(translate_unknown_value)" "]" ) case str(name): diff --git a/tests/codegen/rpc/generated/test_service/rpc_method.py b/tests/codegen/rpc/generated/test_service/rpc_method.py index 2ec649d3..263c36b0 100644 --- a/tests/codegen/rpc/generated/test_service/rpc_method.py +++ b/tests/codegen/rpc/generated/test_service/rpc_method.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError -from replit_river.client import RiverUnknownValue, raise_unknown +from replit_river.client import RiverUnknownValue, translate_unknown_value import replit_river as river diff --git a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py index bc1bc434..426049b9 100644 --- a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py +++ b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnum.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError -from replit_river.client import RiverUnknownValue, raise_unknown +from replit_river.client import RiverUnknownValue, translate_unknown_value import replit_river as river @@ -28,9 +28,9 @@ encode_NeedsenumInput: Callable[["NeedsenumInput"], Any] = lambda x: x NeedsenumOutput = Annotated[ Literal["out_first"] | Literal["out_second"] | RiverUnknownValue, - WrapValidator(raise_unknown), + WrapValidator(translate_unknown_value), ] NeedsenumErrors = Annotated[ Literal["err_first"] | Literal["err_second"] | RiverUnknownValue, - WrapValidator(raise_unknown), + WrapValidator(translate_unknown_value), ] diff --git a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py index 62c3d484..6766eb79 100644 --- a/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py +++ b/tests/codegen/snapshot/snapshots/test_unknown_enum/enumService/needsEnumObject.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError -from replit_river.client import RiverUnknownValue, raise_unknown +from replit_river.client import RiverUnknownValue, translate_unknown_value import replit_river as river @@ -95,7 +95,7 @@ class NeedsenumobjectOutputFooOneOf_out_second(BaseModel): NeedsenumobjectOutputFooOneOf_out_first | NeedsenumobjectOutputFooOneOf_out_second | RiverUnknownValue, - WrapValidator(raise_unknown), + WrapValidator(translate_unknown_value), ] @@ -115,7 +115,7 @@ class NeedsenumobjectErrorsFooAnyOf_1(RiverError): NeedsenumobjectErrorsFooAnyOf_0 | NeedsenumobjectErrorsFooAnyOf_1 | RiverUnknownValue, - WrapValidator(raise_unknown), + WrapValidator(translate_unknown_value), ] diff --git a/tests/codegen/stream/generated/test_service/stream_method.py b/tests/codegen/stream/generated/test_service/stream_method.py index 75bc2817..53267467 100644 --- a/tests/codegen/stream/generated/test_service/stream_method.py +++ b/tests/codegen/stream/generated/test_service/stream_method.py @@ -19,7 +19,7 @@ from pydantic import BaseModel, Field, TypeAdapter, WrapValidator from replit_river.error_schema import RiverError -from replit_river.client import RiverUnknownValue, raise_unknown +from replit_river.client import RiverUnknownValue, translate_unknown_value import replit_river as river