Raise RawConnError in RawConnection

This commit is contained in:
NIC619 2019-09-16 18:37:00 +08:00
parent cb632fa509
commit 4cd5b77f10
No known key found for this signature in database
GPG Key ID: 570C35F5C2D51B17

View File

@ -1,5 +1,6 @@
import asyncio
from .exceptions import RawConnError
from .raw_connection_interface import IRawConnection
@ -23,19 +24,28 @@ class RawConnection(IRawConnection):
self._drain_lock = asyncio.Lock()
async def write(self, data: bytes) -> None:
self.writer.write(data)
try:
self.writer.write(data)
except ConnectionResetError:
raise RawConnError()
# Reference: https://github.com/ethereum/lahja/blob/93610b2eb46969ff1797e0748c7ac2595e130aef/lahja/asyncio/endpoint.py#L99-L102 # noqa: E501
# Use a lock to serialize drain() calls. Circumvents this bug:
# https://bugs.python.org/issue29930
async with self._drain_lock:
await self.writer.drain()
try:
await self.writer.drain()
except ConnectionResetError:
raise RawConnError()
async def read(self, n: int = -1) -> bytes:
"""
Read up to ``n`` bytes from the underlying stream.
This call is delegated directly to the underlying ``self.reader``.
"""
return await self.reader.read(n)
try:
return await self.reader.read(n)
except ConnectionResetError:
raise RawConnError()
async def close(self) -> None:
self.writer.close()