py-libp2p/libp2p/network/stream/net_stream.py

51 lines
1.3 KiB
Python
Raw Normal View History

from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
2019-07-28 22:30:51 +08:00
from .net_stream_interface import INetStream
2018-11-12 05:04:57 +08:00
2018-11-13 02:02:49 +08:00
2018-11-12 05:04:57 +08:00
class NetStream(INetStream):
2019-07-30 15:31:02 +08:00
muxed_stream: IMuxedStream
mplex_conn: IMuxedConn
2019-07-28 14:06:29 +08:00
protocol_id: str
2019-07-30 15:31:02 +08:00
def __init__(self, muxed_stream: IMuxedStream) -> None:
2018-11-12 05:04:57 +08:00
self.muxed_stream = muxed_stream
2019-03-03 20:34:49 +08:00
self.mplex_conn = muxed_stream.mplex_conn
2018-11-13 02:02:49 +08:00
self.protocol_id = None
2018-11-12 05:04:57 +08:00
2019-07-28 14:06:29 +08:00
def get_protocol(self) -> str:
2018-11-12 05:04:57 +08:00
"""
:return: protocol id that stream runs on
"""
return self.protocol_id
2019-07-28 14:06:29 +08:00
def set_protocol(self, protocol_id: str) -> None:
2018-11-12 05:04:57 +08:00
"""
:param protocol_id: protocol id that stream runs on
:return: true if successful
"""
self.protocol_id = protocol_id
2019-07-28 14:06:29 +08:00
async def read(self) -> bytes:
2018-11-12 05:04:57 +08:00
"""
read from stream
:return: bytes of input until EOF
"""
2018-11-13 00:00:43 +08:00
return await self.muxed_stream.read()
2018-11-12 05:04:57 +08:00
2019-07-28 14:06:29 +08:00
async def write(self, data: bytes) -> int:
2018-11-12 05:04:57 +08:00
"""
write to stream
:return: number of bytes written
"""
2018-11-13 02:02:49 +08:00
return await self.muxed_stream.write(data)
2018-11-12 05:04:57 +08:00
2019-07-28 14:06:29 +08:00
async def close(self) -> bool:
2018-11-12 05:04:57 +08:00
"""
close stream
:return: true if successful
"""
2018-11-19 00:22:17 +08:00
await self.muxed_stream.close()
2018-11-12 05:04:57 +08:00
return True