From 22fb9db2bcc9f921836dec95df97e3f3ea2a18d3 Mon Sep 17 00:00:00 2001 From: dolfies Date: Sun, 1 Dec 2024 09:03:10 -0500 Subject: [PATCH] More closing fixes --- curl_cffi/requests/websockets.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/curl_cffi/requests/websockets.py b/curl_cffi/requests/websockets.py index d112122..8c731a4 100644 --- a/curl_cffi/requests/websockets.py +++ b/curl_cffi/requests/websockets.py @@ -381,7 +381,7 @@ def send(self, payload: Union[str, bytes], flags: CurlWsFlag = CurlWsFlag.BINARY flags: flags for the frame. """ if self.closed: - raise TypeError("WebSocket is closed") + raise SessionClosed("WebSocket is closed") # curl expects bytes if isinstance(payload, str): @@ -525,7 +525,7 @@ def loop(self): async def __aiter__(self) -> Self: if self.closed: - raise TypeError("WebSocket is closed") + raise SessionClosed("WebSocket is closed") return self async def __anext__(self) -> bytes: @@ -541,7 +541,7 @@ async def recv_fragment(self, *, timeout: Optional[float] = None) -> Tuple[bytes timeout: how many seconds to wait before giving up. """ if self.closed: - raise TypeError("WebSocket is closed") + raise SessionClosed("WebSocket is closed") if self._recv_lock.locked(): raise TypeError("Concurrent call to recv_fragment() is not allowed") @@ -551,7 +551,7 @@ async def recv_fragment(self, *, timeout: Optional[float] = None) -> Tuple[bytes ) if frame.flags & CurlWsFlag.CLOSE: try: - self._close_code, self._close_reason = self._unpack_close_frame(chunk) + code, message = self._close_code, self._close_reason = self._unpack_close_frame(chunk) except WebSocketError as e: # Follow the spec to close the connection # Errors do not respect autoclose @@ -559,7 +559,7 @@ async def recv_fragment(self, *, timeout: Optional[float] = None) -> Tuple[bytes await self.close(e.code) raise if self.autoclose: - await self.close() + await self.close(code, message.encode()) return chunk, frame @@ -628,7 +628,7 @@ async def send(self, payload: Union[str, bytes], flags: CurlWsFlag = CurlWsFlag. flags: flags for the frame. """ if self.closed: - raise TypeError("WebSocket is closed") + raise SessionClosed("WebSocket is closed") # curl expects bytes if isinstance(payload, str):