codekingpro/portable-devtools

1from __future__ import annotations2 3__all__ = (4    "TextConnectable",5    "TextReceiveStream",6    "TextSendStream",7    "TextStream",8)9 10import codecs11import sys12from collections.abc import Callable, Mapping13from dataclasses import InitVar, dataclass, field14from typing import Any15 16from ..abc import (17    AnyByteReceiveStream,18    AnyByteSendStream,19    AnyByteStream,20    AnyByteStreamConnectable,21    ObjectReceiveStream,22    ObjectSendStream,23    ObjectStream,24    ObjectStreamConnectable,25)26 27if sys.version_info >= (3, 12):28    from typing import override29else:30    from typing_extensions import override31 32 33@dataclass(eq=False)34class TextReceiveStream(ObjectReceiveStream[str]):35    """36    Stream wrapper that decodes bytes to strings using the given encoding.37 38    Decoding is done using :class:`~codecs.IncrementalDecoder` which returns any39    completely received unicode characters as soon as they come in.40 41    :param transport_stream: any bytes-based receive stream42    :param encoding: character encoding to use for decoding bytes to strings (defaults43        to ``utf-8``)44    :param errors: handling scheme for decoding errors (defaults to ``strict``; see the45        `codecs module documentation`_ for a comprehensive list of options)46 47    .. _codecs module documentation:48        https://docs.python.org/3/library/codecs.html#codec-objects49    """50 51    transport_stream: AnyByteReceiveStream52    encoding: InitVar[str] = "utf-8"53    errors: InitVar[str] = "strict"54    _decoder: codecs.IncrementalDecoder = field(init=False)55 56    def __post_init__(self, encoding: str, errors: str) -> None:57        decoder_class = codecs.getincrementaldecoder(encoding)58        self._decoder = decoder_class(errors=errors)59 60    async def receive(self) -> str:61        while True:62            chunk = await self.transport_stream.receive()63            decoded = self._decoder.decode(chunk)64            if decoded:65                return decoded66 67    async def aclose(self) -> None:68        await self.transport_stream.aclose()69        self._decoder.reset()70 71    @property72    def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:73        return self.transport_stream.extra_attributes74 75 76@dataclass(eq=False)77class TextSendStream(ObjectSendStream[str]):78    """79    Sends strings to the wrapped stream as bytes using the given encoding.80 81    :param AnyByteSendStream transport_stream: any bytes-based send stream82    :param str encoding: character encoding to use for encoding strings to bytes83        (defaults to ``utf-8``)84    :param str errors: handling scheme for encoding errors (defaults to ``strict``; see85        the `codecs module documentation`_ for a comprehensive list of options)86 87    .. _codecs module documentation:88        https://docs.python.org/3/library/codecs.html#codec-objects89    """90 91    transport_stream: AnyByteSendStream92    encoding: InitVar[str] = "utf-8"93    errors: str = "strict"94    _encoder: Callable[..., tuple[bytes, int]] = field(init=False)95 96    def __post_init__(self, encoding: str) -> None:97        self._encoder = codecs.getencoder(encoding)98 99    async def send(self, item: str) -> None:100        encoded = self._encoder(item, self.errors)[0]101        await self.transport_stream.send(encoded)102 103    async def aclose(self) -> None:104        await self.transport_stream.aclose()105 106    @property107    def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:108        return self.transport_stream.extra_attributes109 110 111@dataclass(eq=False)112class TextStream(ObjectStream[str]):113    """114    A bidirectional stream that decodes bytes to strings on receive and encodes strings115    to bytes on send.116 117    Extra attributes will be provided from both streams, with the receive stream118    providing the values in case of a conflict.119 120    :param AnyByteStream transport_stream: any bytes-based stream121    :param str encoding: character encoding to use for encoding/decoding strings to/from122        bytes (defaults to ``utf-8``)123    :param str errors: handling scheme for encoding errors (defaults to ``strict``; see124        the `codecs module documentation`_ for a comprehensive list of options)125 126    .. _codecs module documentation:127        https://docs.python.org/3/library/codecs.html#codec-objects128    """129 130    transport_stream: AnyByteStream131    encoding: InitVar[str] = "utf-8"132    errors: InitVar[str] = "strict"133    _receive_stream: TextReceiveStream = field(init=False)134    _send_stream: TextSendStream = field(init=False)135 136    def __post_init__(self, encoding: str, errors: str) -> None:137        self._receive_stream = TextReceiveStream(138            self.transport_stream, encoding=encoding, errors=errors139        )140        self._send_stream = TextSendStream(141            self.transport_stream, encoding=encoding, errors=errors142        )143 144    async def receive(self) -> str:145        return await self._receive_stream.receive()146 147    async def send(self, item: str) -> None:148        await self._send_stream.send(item)149 150    async def send_eof(self) -> None:151        await self.transport_stream.send_eof()152 153    async def aclose(self) -> None:154        await self._send_stream.aclose()155        await self._receive_stream.aclose()156 157    @property158    def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:159        return {160            **self._send_stream.extra_attributes,161            **self._receive_stream.extra_attributes,162        }163 164 165class TextConnectable(ObjectStreamConnectable[str]):166    def __init__(self, connectable: AnyByteStreamConnectable):167        """168        :param connectable: the bytestream endpoint to wrap169 170        """171        self.connectable = connectable172 173    @override174    async def connect(self) -> TextStream:175        stream = await self.connectable.connect()176        return TextStream(stream)177