Optional -> A | None

Author: blast-hardcheese

scripts/parity/gen.py

@@ -1,6 +1,6 @@
 import random
 import string
-from typing import Callable, Optional, TypeVar
+from typing import Callable, TypeVar
 
 A = TypeVar("A")
 
@@ -37,7 +37,7 @@ def gen_choice(choices: list[A]) -> Callable[[], A]:
     return lambda: random.choice(choices)
 
 
-def gen_opt(gen_x: Callable[[], A]) -> Callable[[], Optional[A]]:
+def gen_opt(gen_x: Callable[[], A]) -> Callable[[], A | None]:
     return lambda: gen_x() if gen_bool() else None
 
 

src/replit_river/client.py

@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import timedelta
-from typing import Any, AsyncGenerator, Generator, Generic, Literal, Optional, Union
+from typing import Any, AsyncGenerator, Generator, Generic, Literal, Union
 
 from opentelemetry import trace
 from opentelemetry.trace import Span, SpanKind, Status, StatusCode
@@ -100,9 +100,9 @@ async def send_upload(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
@@ -151,9 +151,9 @@ async def send_stream(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
@@ -186,7 +186,7 @@ class _SpanHandle:
     def set_status(
         self,
         status: Union[Status, StatusCode],
-        description: Optional[str] = None,
+        description: str | None = None,
     ) -> None:
         if self.did_set_status:
             return

src/replit_river/client_session.py

@@ -2,7 +2,7 @@
 import logging
 from collections.abc import AsyncIterable
 from datetime import timedelta
-from typing import Any, AsyncGenerator, Callable, Optional, Union
+from typing import Any, AsyncGenerator, Callable, Union
 
 import nanoid  # type: ignore
 from aiochannel import Channel
@@ -102,9 +102,9 @@ async def send_upload(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],
@@ -241,9 +241,9 @@ async def send_stream(
         self,
         service_name: str,
         procedure_name: str,
-        init: Optional[InitType],
+        init: InitType | None,
         request: AsyncIterable[RequestType],
-        init_serializer: Optional[Callable[[InitType], Any]],
+        init_serializer: Callable[[InitType], Any] | None,
         request_serializer: Callable[[RequestType], Any],
         response_deserializer: Callable[[Any], ResponseType],
         error_deserializer: Callable[[Any], ErrorType],

src/replit_river/client_transport.py

@@ -1,7 +1,7 @@
 import asyncio
 import logging
 from collections.abc import Awaitable, Callable
-from typing import Generic, Optional
+from typing import Generic
 
 import websockets
 from pydantic import ValidationError
@@ -98,7 +98,7 @@ async def get_or_create_session(self) -> ClientSession:
                 await existing_session.close()
                 return await self._create_new_session()
 
-    async def _get_existing_session(self) -> Optional[ClientSession]:
+    async def _get_existing_session(self) -> ClientSession | None:
         async with self._session_lock:
             if not self._sessions:
                 return None
@@ -117,7 +117,7 @@ async def _get_existing_session(self) -> Optional[ClientSession]:
 
     async def _establish_new_connection(
         self,
-        old_session: Optional[ClientSession] = None,
+        old_session: ClientSession | None = None,
     ) -> tuple[
         WebSocketCommonProtocol,
         ControlMessageHandshakeRequest[HandshakeMetadataType],
@@ -129,7 +129,7 @@ async def _establish_new_connection(
         client_id = self._client_id
         logger.info("Attempting to establish new ws connection")
 
-        last_error: Optional[Exception] = None
+        last_error: Exception | None = None
         for i in range(max_retry):
             if i > 0:
                 logger.info(f"Retrying build handshake number {i} times")
@@ -221,7 +221,7 @@ async def _send_handshake_request(
         transport_id: str,
         to_id: str,
         session_id: str,
-        handshake_metadata: Optional[HandshakeMetadataType],
+        handshake_metadata: HandshakeMetadataType | None,
         websocket: WebSocketCommonProtocol,
         expected_session_state: ExpectedSessionState,
     ) -> ControlMessageHandshakeRequest[HandshakeMetadataType]:
@@ -291,7 +291,7 @@ async def _establish_handshake(
         session_id: str,
         handshake_metadata: HandshakeMetadataType,
         websocket: WebSocketCommonProtocol,
-        old_session: Optional[ClientSession],
+        old_session: ClientSession | None,
     ) -> tuple[
         ControlMessageHandshakeRequest[HandshakeMetadataType],
         ControlMessageHandshakeResponse,

src/replit_river/codegen/client.py

@@ -7,7 +7,6 @@
     Any,
     Callable,
     Literal,
-    Optional,
     OrderedDict,
     Sequence,
     Set,
@@ -76,7 +75,6 @@
 from typing import (
     Any,
     Literal,
-    Optional,
     Mapping,
     NotRequired,
     Union,
@@ -96,11 +94,11 @@
 
 
 class RiverConcreteType(BaseModel):
-    type: Optional[str] = Field(default=None)
+    type: str | None = Field(default=None)
     properties: dict[str, "RiverType"] = Field(default_factory=lambda: dict())
     required: Set[str] = Field(default=set())
-    items: Optional["RiverType"] = Field(default=None)
-    const: Optional[Union[str, int]] = Field(default=None)
+    items: "RiverType | None" = Field(default=None)
+    const: Union[str, int] | None = Field(default=None)
     patternProperties: dict[str, "RiverType"] = Field(default_factory=lambda: dict())
 
 
@@ -124,14 +122,14 @@ class RiverNotType(BaseModel):
 
 
 class RiverProcedure(BaseModel):
-    init: Optional[RiverType] = Field(default=None)
+    init: RiverType | None = Field(default=None)
     input: RiverType
     output: RiverType
-    errors: Optional[RiverType] = Field(default=None)
+    errors: RiverType | None = Field(default=None)
     type: (
         Literal["rpc"] | Literal["stream"] | Literal["subscription"] | Literal["upload"]
     )
-    description: Optional[str] = Field(default=None)
+    description: str | None = Field(default=None)
 
 
 class RiverService(BaseModel):
@@ -140,7 +138,7 @@ class RiverService(BaseModel):
 
 class RiverSchema(BaseModel):
     services: dict[str, RiverService]
-    handshakeSchema: Optional[RiverConcreteType] = Field(default=None)
+    handshakeSchema: RiverConcreteType | None = Field(default=None)
 
 
 RiverSchemaFile = RootModel[RiverSchema]
@@ -640,7 +638,7 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
                                 """
                             )
                         current_chunks.append(
-                            f"  kind: Optional[{render_type_expr(type_name)}]{value}"
+                            f"  kind: {render_type_expr(type_name)} | None{value}"
                         )
                     else:
                         value = ""
@@ -663,7 +661,7 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
                                 reindent(
                                     "  ",
                                     f"""\
-                        {name}: NotRequired[Optional[{render_type_expr(type_name)}]]
+                        {name}: NotRequired[{render_type_expr(type_name)}] | None
                                 """,
                                 )
                             )
@@ -672,7 +670,7 @@ def extract_props(tpe: RiverType) -> list[dict[str, RiverType]]:
                                 reindent(
                                     "  ",
                                     f"""\
-                        {name}: Optional[{render_type_expr(type_name)}] = None
+                        {name}: {render_type_expr(type_name)} | None = None
                                 """,
                                 )
                             )
@@ -772,7 +770,7 @@ def __init__(self, client: river.Client[Any]):
     ]
     for name, procedure in schema.procedures.items():
         module_names = [ModuleName(name)]
-        init_type: Optional[TypeExpression] = None
+        init_type: TypeExpression | None = None
         if procedure.init:
             init_type, module_info, init_chunks, encoder_names = encode_type(
                 procedure.init,
@@ -852,7 +850,7 @@ def __init__(self, client: river.Client[Any]):
                             """
 
         # Init renderer
-        render_init_method: Optional[str] = None
+        render_init_method: str | None = None
         if init_type and procedure.init is not None:
             if input_base_class == "TypedDict":
                 if is_literal(procedure.init):
@@ -878,7 +876,7 @@ def __init__(self, client: river.Client[Any]):
         )
 
         # Input renderer
-        render_input_method: Optional[str] = None
+        render_input_method: str | None = None
         if input_base_class == "TypedDict":
             if is_literal(procedure.input):
                 render_input_method = "lambda x: x"

src/replit_river/error_schema.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -38,7 +38,7 @@ class RiverServiceException(RiverException):
     """Exception raised by river as a result of a fault in the service running river."""
 
     def __init__(
-        self, code: str, message: str, service: Optional[str], procedure: Optional[str]
+        self, code: str, message: str, service: str | None, procedure: str | None
     ) -> None:
         self.code = code
         self.message = message
@@ -92,7 +92,7 @@ def stringify_exception(e: BaseException, limit: int = 10) -> str:
         # If there are no causes, just fall back to stringifying the exception.
         return str(e)
     causes: list[str] = []
-    cause: Optional[BaseException] = e
+    cause: BaseException | None = e
     while cause and limit:
         causes.append(str(cause))
         cause = cause.__cause__

src/replit_river/message_buffer.py

@@ -1,6 +1,5 @@
 import asyncio
 import logging
-from typing import Optional
 
 from replit_river.rpc import TransportMessage
 from replit_river.transport_options import MAX_MESSAGE_BUFFER_SIZE
@@ -41,7 +40,7 @@ async def put(self, message: TransportMessage) -> None:
                 raise MessageBufferClosedError("message buffer is closed")
             self.buffer.append(message)
 
-    async def peek(self) -> Optional[TransportMessage]:
+    async def peek(self) -> TransportMessage | None:
         """Peek the first message in the buffer, returns None if the buffer is empty."""
         async with self._lock:
             if len(self.buffer) == 0:

src/replit_river/rpc.py

@@ -12,7 +12,6 @@
     Literal,
     Mapping,
     NoReturn,
-    Optional,
     Sequence,
     TypeVar,
     Union,
@@ -63,22 +62,22 @@
 # Equivalent of https://github.com/replit/river/blob/c1345f1ff6a17a841d4319fad5c153b5bda43827/transport/message.ts#L23-L33
 class ExpectedSessionState(BaseModel):
     nextExpectedSeq: int
-    nextSentSeq: Optional[int] = None
+    nextSentSeq: int | None = None
 
 
 class ControlMessageHandshakeRequest(BaseModel, Generic[HandshakeMetadataType]):
     type: Literal["HANDSHAKE_REQ"] = "HANDSHAKE_REQ"
     protocolVersion: str
     sessionId: str
     expectedSessionState: ExpectedSessionState
-    metadata: Optional[HandshakeMetadataType] = None
+    metadata: HandshakeMetadataType | None = None
 
 
 class HandShakeStatus(BaseModel):
     ok: bool
-    sessionId: Optional[str] = None
-    reason: Optional[str] = None
-    code: Optional[str] = None
+    sessionId: str | None = None
+    reason: str | None = None
+    code: str | None = None
 
 
 class ControlMessageHandshakeResponse(BaseModel):
@@ -98,11 +97,11 @@ class TransportMessage(BaseModel):
     to: str
     seq: int
     ack: int
-    serviceName: Optional[str] = None
-    procedureName: Optional[str] = None
+    serviceName: str | None = None
+    procedureName: str | None = None
     streamId: str
     controlFlags: int
-    tracing: Optional[PropagationContext] = None
+    tracing: PropagationContext | None = None
     payload: Any
     model_config = ConfigDict(populate_by_name=True)
     # need this because we create TransportMessage objects with destructuring
@@ -131,8 +130,8 @@ class GrpcContext(grpc.aio.ServicerContext, Generic[RequestType, ResponseType]):
 
     def __init__(self, peer: str) -> None:
         self._peer = peer
-        self._abort_code: Optional[grpc.StatusCode] = None
-        self._abort_details: Optional[str] = None
+        self._abort_code: grpc.StatusCode | None = None
+        self._abort_details: str | None = None
 
     async def abort(
         self,
@@ -157,10 +156,10 @@ def invocation_metadata(self) -> None:
     def peer(self) -> str:
         return self._peer
 
-    def peer_identities(self) -> Optional[Iterable[bytes]]:
+    def peer_identities(self) -> Iterable[bytes] | None:
         return None
 
-    def peer_identity_key(self) -> Optional[str]:
+    def peer_identity_key(self) -> str | None:
         return None
 
     async def read(self) -> RequestType:

src/replit_river/server.py

@@ -1,6 +1,6 @@
 import asyncio
 import logging
-from typing import Mapping, Optional
+from typing import Mapping
 
 import websockets
 from websockets.exceptions import ConnectionClosed
@@ -41,7 +41,7 @@ def add_rpc_handlers(
 
     async def _handshake_to_get_session(
         self, websocket: WebSocketServerProtocol
-    ) -> Optional[Session]:
+    ) -> Session | None:
         """This is a wrapper to make sentry happy, sentry doesn't recognize the
         exception handling outside of a task or asyncio.wait_for. So we need to catch
         the errors specifically here.