33 lines
992 B
Python
33 lines
992 B
Python
|
import trio
|
||
|
|
from trio import SocketStream
|
||
|
|
from libp2p.io.abc import ReadWriteCloser
|
||
|
|
from libp2p.io.exceptions import IOException
|
||
|
|
import logging
|
||
|
|
|
||
|
|
|
||
|
|
logger = logging.getLogger("libp2p.io.trio")
|
||
|
|
|
||
|
|
|
||
|
|
class TrioReadWriteCloser(ReadWriteCloser):
|
||
|
|
stream: SocketStream
|
||
|
|
|
||
|
|
def __init__(self, stream: SocketStream) -> None:
|
||
|
|
self.stream = stream
|
||
|
|
|
||
|
|
async def write(self, data: bytes) -> None:
|
||
|
|
"""Raise `RawConnError` if the underlying connection breaks."""
|
||
|
|
try:
|
||
|
|
await self.stream.send_all(data)
|
||
|
|
except (trio.ClosedResourceError, trio.BrokenResourceError) as error:
|
||
|
|
raise IOException(error)
|
||
|
|
|
||
|
|
async def read(self, n: int = -1) -> bytes:
|
||
|
|
max_bytes = n if n != -1 else None
|
||
|
|
try:
|
||
|
|
return await self.stream.receive_some(max_bytes)
|
||
|
|
except (trio.ClosedResourceError, trio.BrokenResourceError) as error:
|
||
|
|
raise IOException(error)
|
||
|
|
|
||
|
|
async def close(self) -> None:
|
||
|
|
await self.stream.aclose()
|