Stricter error types (#141)

Why
===

Have the generated code return error types that actually satisfy the
`ErrorType`. Previously the union type for error types generated can be
values of `RiverUnknownValue` which is not an error type. This
introduces a `RiverUnknownError` and uses that in the error unions
instead. Also type adapter types are now more strictly typed.

# Testing

When switching the client code over, we should get errors about
`RiverUnknownError` not being handled and force you to fix them.

---------

Co-authored-by: Devon Stewart <[email protected]>

Author: airportyh

src/replit_river/client.py

@@ -13,7 +13,7 @@
 )
 
 from replit_river.client_transport import ClientTransport
-from replit_river.error_schema import RiverError, RiverException
+from replit_river.error_schema import ERROR_CODE_UNKNOWN, RiverError, RiverException
 from replit_river.transport_options import (
     HandshakeMetadataType,
     TransportOptions,
@@ -37,6 +37,10 @@ class RiverUnknownValue(BaseModel):
     value: Any
 
 
+class RiverUnknownError(RiverError):
+    pass
+
+
 def translate_unknown_value(
     value: Any, handler: Callable[[Any], Any], info: ValidationInfo
 ) -> Any | RiverUnknownValue:
@@ -46,6 +50,21 @@ def translate_unknown_value(
         return RiverUnknownValue(tag="RiverUnknownValue", value=value)
 
 
+def translate_unknown_error(
+    value: Any, handler: Callable[[Any], Any], info: ValidationInfo
+) -> Any | RiverUnknownError:
+    try:
+        return handler(value)
+    except Exception:
+        if isinstance(value, dict) and "code" in value and "message" in value:
+            return RiverUnknownError(
+                code=value["code"],
+                message=value["message"],
+            )
+        else:
+            return RiverUnknownError(code=ERROR_CODE_UNKNOWN, message="Unknown error")
+
+
 class Client(Generic[HandshakeMetadataType]):
     def __init__(
         self,

src/replit_river/codegen/client.py

@@ -81,7 +81,8 @@
 
 from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
 from replit_river.error_schema import RiverError
-from replit_river.client import RiverUnknownValue, translate_unknown_value
+from replit_river.client import RiverUnknownError, translate_unknown_error, \
+    RiverUnknownValue, translate_unknown_value
 
 import replit_river as river
 
@@ -154,6 +155,20 @@ def encode_type(
     in_module: list[ModuleName],
     permit_unknown_members: bool,
 ) -> tuple[TypeExpression, list[ModuleName], list[FileContents], set[TypeName]]:
+    def _make_open_union_type_expr(one_of: list[TypeExpression]) -> OpenUnionTypeExpr:
+        if base_model == "RiverError":
+            return OpenUnionTypeExpr(
+                UnionTypeExpr(one_of),
+                fallback_type="RiverUnknownError",
+                validator_function="translate_unknown_error",
+            )
+        else:
+            return OpenUnionTypeExpr(
+                UnionTypeExpr(one_of),
+                fallback_type="RiverUnknownValue",
+                validator_function="translate_unknown_value",
+            )
+
     encoder_name: TypeName | None = None  # defining this up here to placate mypy
     chunks: list[FileContents] = []
     if isinstance(type, RiverNotType):
@@ -304,7 +319,7 @@ def flatten_union(tpe: RiverType) -> list[RiverType]:
                     )
                 union: TypeExpression
                 if permit_unknown_members:
-                    union = OpenUnionTypeExpr(UnionTypeExpr(one_of))
+                    union = _make_open_union_type_expr(one_of)
                 else:
                     union = UnionTypeExpr(one_of)
                 chunks.append(
@@ -383,7 +398,7 @@ def {_field_name}(
                             )
                             raise ValueError(f"What does it mean to have {_o2} here?")
         if permit_unknown_members:
-            union = OpenUnionTypeExpr(UnionTypeExpr(any_of))
+            union = _make_open_union_type_expr(any_of)
         else:
             union = UnionTypeExpr(any_of)
         if is_literal(type):
@@ -795,17 +810,18 @@ def _type_adapter_definition(
         _type: TypeExpression,
         module_info: list[ModuleName],
     ) -> tuple[list[TypeName], list[ModuleName], list[FileContents]]:
+        varname = render_type_expr(type_adapter_name)
         rendered_type_expr = render_type_expr(_type)
         return (
             [type_adapter_name],
             module_info,
             [
                 FileContents(
                     dedent(f"""
-                    {render_type_expr(type_adapter_name)}: TypeAdapter[Any] = (
-                        TypeAdapter({rendered_type_expr})
-                    )
-                """)
+                {varname}: TypeAdapter[{rendered_type_expr}] = (
+                    TypeAdapter({rendered_type_expr})
+                )
+            """)
                 )
             ],
         )

src/replit_river/codegen/typing.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import NewType, assert_never
+from typing import NewType, assert_never, cast
 
 ModuleName = NewType("ModuleName", str)
 ClassName = NewType("ClassName", str)
@@ -96,6 +96,8 @@ def __lt__(self, other: object) -> bool:
 @dataclass(frozen=True)
 class OpenUnionTypeExpr:
     union: UnionTypeExpr
+    fallback_type: str
+    validator_function: str
 
     def __str__(self) -> str:
         raise Exception("Complex type must be put through render_type_expr!")
@@ -182,10 +184,11 @@ def render_type_expr(value: TypeExpression) -> str:
                     retval = "None"
             return retval
         case OpenUnionTypeExpr(inner):
+            open_union = cast(OpenUnionTypeExpr, value)
             return (
                 "Annotated["
-                f"{render_type_expr(inner)} | RiverUnknownValue,"
-                "WrapValidator(translate_unknown_value)"
+                f"{render_type_expr(inner)} | {open_union.fallback_type},"
+                f"WrapValidator({open_union.validator_function})"
                 "]"
             )
         case TypeName(name):

src/replit_river/error_schema.py

@@ -17,6 +17,9 @@
 # ERROR_CODE_CANCEL is the code used when either server or client cancels the stream.
 ERROR_CODE_CANCEL = "CANCEL"
 
+# ERROR_CODE_UNKNOWN is the code for the RiverUnknownError
+ERROR_CODE_UNKNOWN = "UNKNOWN"
+
 
 class RiverError(BaseModel):
     """Error message from the server."""

tests/codegen/rpc/generated/test_service/rpc_method.py

@@ -12,7 +12,12 @@
 
 from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
 from replit_river.error_schema import RiverError
-from replit_river.client import RiverUnknownValue, translate_unknown_value
+from replit_river.client import (
+    RiverUnknownError,
+    translate_unknown_error,
+    RiverUnknownValue,
+    translate_unknown_value,
+)
 
 import replit_river as river
 
@@ -35,11 +40,13 @@ class Rpc_MethodInput(TypedDict):
     data: str
 
 
-Rpc_MethodInputTypeAdapter: TypeAdapter[Any] = TypeAdapter(Rpc_MethodInput)
+Rpc_MethodInputTypeAdapter: TypeAdapter[Rpc_MethodInput] = TypeAdapter(Rpc_MethodInput)
 
 
 class Rpc_MethodOutput(BaseModel):
     data: str
 
 
-Rpc_MethodOutputTypeAdapter: TypeAdapter[Any] = TypeAdapter(Rpc_MethodOutput)
+Rpc_MethodOutputTypeAdapter: TypeAdapter[Rpc_MethodOutput] = TypeAdapter(
+    Rpc_MethodOutput
+)

tests/codegen/snapshot/snapshots/test_basic_stream/test_service/__init__.py

@@ -16,6 +16,12 @@
     Stream_MethodOutputTypeAdapter,
     encode_Stream_MethodInput,
 )
+from .emit_error import Emit_ErrorErrors, Emit_ErrorErrorsTypeAdapter
+
+intTypeAdapter: TypeAdapter[int] = TypeAdapter(int)
+
+
+boolTypeAdapter: TypeAdapter[bool] = TypeAdapter(bool)
 
 
 class Test_ServiceService:
@@ -40,3 +46,22 @@ async def stream_method(
                 x  # type: ignore[arg-type]
             ),
         )
+
+    async def emit_error(
+        self,
+        inputStream: AsyncIterable[int],
+    ) -> AsyncIterator[bool | Emit_ErrorErrors | RiverError]:
+        return self.client.send_stream(
+            "test_service",
+            "emit_error",
+            None,
+            inputStream,
+            None,
+            lambda x: x,
+            lambda x: boolTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+            lambda x: Emit_ErrorErrorsTypeAdapter.validate_python(
+                x  # type: ignore[arg-type]
+            ),
+        )

tests/codegen/snapshot/snapshots/test_basic_stream/test_service/emit_error.py

@@ -0,0 +1,45 @@
+# 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
+
+
+class Emit_ErrorErrorsOneOf_DATA_LOSS(RiverError):
+    code: Literal["DATA_LOSS"]
+    message: str
+
+
+class Emit_ErrorErrorsOneOf_UNEXPECTED_DISCONNECT(RiverError):
+    code: Literal["UNEXPECTED_DISCONNECT"]
+    message: str
+
+
+Emit_ErrorErrors = Annotated[
+    Emit_ErrorErrorsOneOf_DATA_LOSS
+    | Emit_ErrorErrorsOneOf_UNEXPECTED_DISCONNECT
+    | RiverUnknownError,
+    WrapValidator(translate_unknown_error),
+]
+
+
+Emit_ErrorErrorsTypeAdapter: TypeAdapter[Emit_ErrorErrors] = TypeAdapter(
+    Emit_ErrorErrors
+)

tests/codegen/snapshot/snapshots/test_basic_stream/test_service/stream_method.py

@@ -12,7 +12,12 @@
 
 from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
 from replit_river.error_schema import RiverError
-from replit_river.client import RiverUnknownValue, translate_unknown_value
+from replit_river.client import (
+    RiverUnknownError,
+    translate_unknown_error,
+    RiverUnknownValue,
+    translate_unknown_value,
+)
 
 import replit_river as river
 
@@ -35,11 +40,15 @@ class Stream_MethodInput(TypedDict):
     data: str
 
 
-Stream_MethodInputTypeAdapter: TypeAdapter[Any] = TypeAdapter(Stream_MethodInput)
+Stream_MethodInputTypeAdapter: TypeAdapter[Stream_MethodInput] = TypeAdapter(
+    Stream_MethodInput
+)
 
 
 class Stream_MethodOutput(BaseModel):
     data: str
 
 
-Stream_MethodOutputTypeAdapter: TypeAdapter[Any] = TypeAdapter(Stream_MethodOutput)
+Stream_MethodOutputTypeAdapter: TypeAdapter[Stream_MethodOutput] = TypeAdapter(
+    Stream_MethodOutput
+)

tests/codegen/snapshot/snapshots/test_pathological_types/test_service/__init__.py

@@ -31,7 +31,7 @@
     encode_Pathological_MethodInputReq_Obj_Undefined,
 )
 
-boolTypeAdapter: TypeAdapter[Any] = TypeAdapter(bool)
+boolTypeAdapter: TypeAdapter[bool] = TypeAdapter(bool)
 
 
 class Test_ServiceService:

tests/codegen/snapshot/snapshots/test_pathological_types/test_service/pathological_method.py

@@ -12,7 +12,12 @@
 
 from pydantic import BaseModel, Field, TypeAdapter, WrapValidator
 from replit_river.error_schema import RiverError
-from replit_river.client import RiverUnknownValue, translate_unknown_value
+from replit_river.client import (
+    RiverUnknownError,
+    translate_unknown_error,
+    RiverUnknownValue,
+    translate_unknown_value,
+)
 
 import replit_river as river
 
@@ -473,6 +478,6 @@ class Pathological_MethodInput(TypedDict):
     undefined: NotRequired[None]
 
 
-Pathological_MethodInputTypeAdapter: TypeAdapter[Any] = TypeAdapter(
-    Pathological_MethodInput
+Pathological_MethodInputTypeAdapter: TypeAdapter[Pathological_MethodInput] = (
+    TypeAdapter(Pathological_MethodInput)
 )