1  from __future__ import annotations
   2  
   3  import codecs
   4  from collections.abc import Callable, Mapping
   5  from dataclasses import InitVar, dataclass, field
   6  from typing import Any
   7  
   8  from ..abc import (
   9      AnyByteReceiveStream,
  10      AnyByteSendStream,
  11      AnyByteStream,
  12      ObjectReceiveStream,
  13      ObjectSendStream,
  14      ObjectStream,
  15  )
  16  
  17  
  18  @dataclass(eq=False)
  19  class TextReceiveStream(ObjectReceiveStream[str]):
  20      """
  21      Stream wrapper that decodes bytes to strings using the given encoding.
  22  
  23      Decoding is done using :class:`~codecs.IncrementalDecoder` which returns any
  24      completely received unicode characters as soon as they come in.
  25  
  26      :param transport_stream: any bytes-based receive stream
  27      :param encoding: character encoding to use for decoding bytes to strings (defaults
  28          to ``utf-8``)
  29      :param errors: handling scheme for decoding errors (defaults to ``strict``; see the
  30          `codecs module documentation`_ for a comprehensive list of options)
  31  
  32      .. _codecs module documentation:
  33          https://docs.python.org/3/library/codecs.html#codec-objects
  34      """
  35  
  36      transport_stream: AnyByteReceiveStream
  37      encoding: InitVar[str] = "utf-8"
  38      errors: InitVar[str] = "strict"
  39      _decoder: codecs.IncrementalDecoder = field(init=False)
  40  
  41      def __post_init__(self, encoding: str, errors: str) -> None:
  42          decoder_class = codecs.getincrementaldecoder(encoding)
  43          self._decoder = decoder_class(errors=errors)
  44  
  45      async def receive(self) -> str:
  46          while True:
  47              chunk = await self.transport_stream.receive()
  48              decoded = self._decoder.decode(chunk)
  49              if decoded:
  50                  return decoded
  51  
  52      async def aclose(self) -> None:
  53          await self.transport_stream.aclose()
  54          self._decoder.reset()
  55  
  56      @property
  57      def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  58          return self.transport_stream.extra_attributes
  59  
  60  
  61  @dataclass(eq=False)
  62  class TextSendStream(ObjectSendStream[str]):
  63      """
  64      Sends strings to the wrapped stream as bytes using the given encoding.
  65  
  66      :param AnyByteSendStream transport_stream: any bytes-based send stream
  67      :param str encoding: character encoding to use for encoding strings to bytes
  68          (defaults to ``utf-8``)
  69      :param str errors: handling scheme for encoding errors (defaults to ``strict``; see
  70          the `codecs module documentation`_ for a comprehensive list of options)
  71  
  72      .. _codecs module documentation:
  73          https://docs.python.org/3/library/codecs.html#codec-objects
  74      """
  75  
  76      transport_stream: AnyByteSendStream
  77      encoding: InitVar[str] = "utf-8"
  78      errors: str = "strict"
  79      _encoder: Callable[..., tuple[bytes, int]] = field(init=False)
  80  
  81      def __post_init__(self, encoding: str) -> None:
  82          self._encoder = codecs.getencoder(encoding)
  83  
  84      async def send(self, item: str) -> None:
  85          encoded = self._encoder(item, self.errors)[0]
  86          await self.transport_stream.send(encoded)
  87  
  88      async def aclose(self) -> None:
  89          await self.transport_stream.aclose()
  90  
  91      @property
  92      def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
  93          return self.transport_stream.extra_attributes
  94  
  95  
  96  @dataclass(eq=False)
  97  class TextStream(ObjectStream[str]):
  98      """
  99      A bidirectional stream that decodes bytes to strings on receive and encodes strings
 100      to bytes on send.
 101  
 102      Extra attributes will be provided from both streams, with the receive stream
 103      providing the values in case of a conflict.
 104  
 105      :param AnyByteStream transport_stream: any bytes-based stream
 106      :param str encoding: character encoding to use for encoding/decoding strings to/from
 107          bytes (defaults to ``utf-8``)
 108      :param str errors: handling scheme for encoding errors (defaults to ``strict``; see
 109          the `codecs module documentation`_ for a comprehensive list of options)
 110  
 111      .. _codecs module documentation:
 112          https://docs.python.org/3/library/codecs.html#codec-objects
 113      """
 114  
 115      transport_stream: AnyByteStream
 116      encoding: InitVar[str] = "utf-8"
 117      errors: InitVar[str] = "strict"
 118      _receive_stream: TextReceiveStream = field(init=False)
 119      _send_stream: TextSendStream = field(init=False)
 120  
 121      def __post_init__(self, encoding: str, errors: str) -> None:
 122          self._receive_stream = TextReceiveStream(
 123              self.transport_stream, encoding=encoding, errors=errors
 124          )
 125          self._send_stream = TextSendStream(
 126              self.transport_stream, encoding=encoding, errors=errors
 127          )
 128  
 129      async def receive(self) -> str:
 130          return await self._receive_stream.receive()
 131  
 132      async def send(self, item: str) -> None:
 133          await self._send_stream.send(item)
 134  
 135      async def send_eof(self) -> None:
 136          await self.transport_stream.send_eof()
 137  
 138      async def aclose(self) -> None:
 139          await self._send_stream.aclose()
 140          await self._receive_stream.aclose()
 141  
 142      @property
 143      def extra_attributes(self) -> Mapping[Any, Callable[[], Any]]:
 144          return {
 145              **self._send_stream.extra_attributes,
 146              **self._receive_stream.extra_attributes,
 147          }
 148