Skip to content

Commit

Permalink
Fix client API version bugs
Browse files Browse the repository at this point in the history
  • Loading branch information
brentyi committed Aug 2, 2023
1 parent f201a56 commit 1faea9d
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions viser/infra/_infra.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@
Callable,
Dict,
List,
Literal,
NewType,
Optional,
Sequence,
Expand All @@ -32,6 +31,7 @@
from rich import box, style
from rich.panel import Panel
from rich.table import Table
from typing_extensions import Literal, assert_never
from websockets.legacy.server import WebSocketServerProtocol

from ._async_message_buffer import AsyncMessageBuffer, MessageWindow
Expand Down Expand Up @@ -134,7 +134,7 @@ def __init__(
self._message_class = message_class
self._http_server_root = http_server_root
self._verbose = verbose
self._client_api_version = client_api_version
self._client_api_version: Literal[0, 1] = client_api_version

self._thread_executor = ThreadPoolExecutor(max_workers=32)

Expand Down Expand Up @@ -332,7 +332,7 @@ async def _client_producer(
websocket: WebSocketServerProtocol,
client_id: ClientId,
get_next: Callable[[], Awaitable[Message]],
client_api_version: int,
client_api_version: Literal[0, 1],
) -> None:
"""Infinite loop to send messages from a buffer to a single client."""

Expand All @@ -343,34 +343,38 @@ async def _client_producer(
message_future = asyncio.ensure_future(get_next())
outgoing = window.get_window_to_send()
if outgoing is not None:
if client_api_version:
if client_api_version == 1:
serialized = msgpack.packb(
tuple(message.as_serializable_dict() for message in outgoing)
)
assert isinstance(serialized, bytes)
await websocket.send(serialized)
else:
elif client_api_version == 0:
for msg in outgoing:
await websocket.send(msg.as_serializable_dict())
else:
assert_never(client_api_version)


async def _broadcast_producer(
websocket: WebSocketServerProtocol,
get_next_window: Callable[[], Awaitable[Sequence[Message]]],
client_api_version: int,
client_api_version: Literal[0, 1],
) -> None:
"""Infinite loop to broadcast windows of messages from a buffer."""
while True:
outgoing = await get_next_window()
if client_api_version:
if client_api_version == 1:
serialized = msgpack.packb(
tuple(message.as_serializable_dict() for message in outgoing)
)
assert isinstance(serialized, bytes)
await websocket.send(serialized)
else:
elif client_api_version == 0:
for msg in outgoing:
await websocket.send(msg.as_serializable_dict())
else:
assert_never(client_api_version)


async def _consumer(
Expand Down

0 comments on commit 1faea9d

Please sign in to comment.