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

73 lines
2.3 KiB
Python
Raw Normal View History

from typing import Optional
from libp2p.stream_muxer.abc import IMuxedStream
2019-09-09 15:45:35 +08:00
from libp2p.stream_muxer.exceptions import (
MuxedStreamClosed,
MuxedStreamEOF,
MuxedStreamReset,
)
from libp2p.typing import TProtocol
2019-07-28 22:30:51 +08:00
2019-09-09 15:45:35 +08:00
from .exceptions import StreamClosed, StreamEOF, StreamReset
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
2019-09-09 15:45:35 +08:00
# TODO: Handle exceptions from `muxed_stream`
# TODO: Add stream state
# - Reference: https://github.com/libp2p/go-libp2p-swarm/blob/99831444e78c8f23c9335c17d8f7c700ba25ca14/swarm_stream.go # noqa: E501
2018-11-12 05:04:57 +08:00
class NetStream(INetStream):
2019-07-30 15:31:02 +08:00
muxed_stream: IMuxedStream
protocol_id: Optional[TProtocol]
2019-07-28 14:06:29 +08:00
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
self.muxed_conn = muxed_stream.muxed_conn
2018-11-13 02:02:49 +08:00
self.protocol_id = None
2018-11-12 05:04:57 +08:00
def get_protocol(self) -> TProtocol:
2018-11-12 05:04:57 +08:00
"""
:return: protocol id that stream runs on
"""
return self.protocol_id
def set_protocol(self, protocol_id: TProtocol) -> None:
2018-11-12 05:04:57 +08:00
"""
:param protocol_id: protocol id that stream runs on
"""
self.protocol_id = protocol_id
2019-08-07 15:23:20 +08:00
async def read(self, n: int = -1) -> bytes:
"""reads from stream.
2019-08-07 15:23:20 +08:00
:param n: number of bytes to read
:return: bytes of input
2018-11-12 05:04:57 +08:00
"""
2019-09-09 15:45:35 +08:00
try:
return await self.muxed_stream.read(n)
except MuxedStreamEOF as error:
raise StreamEOF from error
except MuxedStreamReset as error:
raise StreamReset from error
2018-11-12 05:04:57 +08:00
2019-07-28 14:06:29 +08:00
async def write(self, data: bytes) -> int:
"""write to stream.
2018-11-12 05:04:57 +08:00
:return: number of bytes written
"""
2019-09-09 15:45:35 +08:00
try:
return await self.muxed_stream.write(data)
except MuxedStreamClosed as error:
raise StreamClosed from error
2018-11-12 05:04:57 +08:00
async def close(self) -> None:
"""close stream."""
2018-11-19 00:22:17 +08:00
await self.muxed_stream.close()
2019-08-27 02:38:39 +08:00
2019-09-09 15:45:35 +08:00
async def reset(self) -> None:
await self.muxed_stream.reset()
2019-09-17 23:38:11 +08:00
# TODO: `remove`: Called by close and write when the stream is in specific states.
2019-09-23 15:01:58 +08:00
# It notifies `ClosedStream` after `SwarmConn.remove_stream` is called.
2019-09-17 23:38:11 +08:00
# Reference: https://github.com/libp2p/go-libp2p-swarm/blob/99831444e78c8f23c9335c17d8f7c700ba25ca14/swarm_stream.go # noqa: E501