From 004c92edbe0c290945def496b43fc4ea821e5331 Mon Sep 17 00:00:00 2001 From: Andrey Egorov Date: Sun, 19 Nov 2017 13:04:40 +0300 Subject: [PATCH] SocketClosed exception --- telethon/extensions/tcp_client.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/telethon/extensions/tcp_client.py b/telethon/extensions/tcp_client.py index 23268e26..1a7be86f 100644 --- a/telethon/extensions/tcp_client.py +++ b/telethon/extensions/tcp_client.py @@ -18,6 +18,9 @@ CONN_RESET_ERRNOS = { class TcpClient: + class SocketClosed(ConnectionError): + pass + def __init__(self, proxy=None, timeout=timedelta(seconds=5), loop=None): self.proxy = proxy self._socket = None @@ -100,7 +103,7 @@ class TcpClient: loop=self._loop ) if not self.connected: - raise ConnectionResetError('Socket has closed') + raise self.SocketClosed() if not done: raise TimeoutError() return await done.pop() @@ -111,6 +114,8 @@ class TcpClient: raise ConnectionResetError('No connection') try: await self._wait_close(self.sock_sendall(data)) + except self.SocketClosed: + raise ConnectionResetError('Socket has closed') except OSError as e: if e.errno in CONN_RESET_ERRNOS: self._raise_connection_reset(e) @@ -128,6 +133,8 @@ class TcpClient: raise ConnectionResetError('No connection') try: partial = await self._wait_close(self.sock_recv(bytes_left)) + except self.SocketClosed: + raise ConnectionResetError('Socket has closed') except OSError as e: if e.errno in CONN_RESET_ERRNOS: self._raise_connection_reset(e)