Simplify frame handling in `ws`
@@ -15,7 +15,6 @@ from tornado.iostream import StreamClosedError
from tornado.websocket import WebSocketClosedError, WebSocketHandler, websocket_connect
import dask
-from dask.utils import ensure_bytes
from distributed.comm.addressing import parse_host_port, unparse_host_port
from distributed.comm.core import (
@@ -37,7 +36,6 @@ from distributed.comm.utils import (
get_tcp_server_address,
to_frames,
)
-from distributed.utils import nbytes
logger = logging.getLogger(__name__)
@@ -138,14 +136,18 @@ class WSHandlerComm(Comm):
frame_split_size=BIG_BYTES_SHARD_SIZE,
)
n = struct.pack("Q", len(frames))
+ nbytes_frames = 0
try:
await self.handler.write_message(n, binary=True)
for frame in frames:
- await self.handler.write_message(ensure_bytes(frame), binary=True)
+ if type(frame) is not bytes:
+ frame = bytes(frame)
+ await self.handler.write_message(frame, binary=True)
+ nbytes_frames += len(frame)
except WebSocketClosedError as e:
raise CommClosedError(str(e))
- return sum(map(nbytes, frames))
+ return nbytes_frames
def abort(self):
self.handler.close()
@@ -227,14 +229,18 @@ class WS(Comm):
frame_split_size=BIG_BYTES_SHARD_SIZE,
)
n = struct.pack("Q", len(frames))
+ nbytes_frames = 0
try:
await self.sock.write_message(n, binary=True)
for frame in frames:
- await self.sock.write_message(ensure_bytes(frame), binary=True)
+ if type(frame) is not bytes:
+ frame = bytes(frame)
+ await self.sock.write_message(frame, binary=True)
+ nbytes_frames += len(frame)
except WebSocketClosedError as e:
raise CommClosedError(e)
- return sum(map(nbytes, frames))
+ return nbytes_frames
async def close(self):
if not self.sock.close_code:
