streams.py 2.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. from typing import Generator
  2. class StreamReader:
  3. """
  4. Generator-based stream reader.
  5. This class doesn't support concurrent calls to :meth:`read_line()`,
  6. :meth:`read_exact()`, or :meth:`read_to_eof()`. Make sure calls are
  7. serialized.
  8. """
  9. def __init__(self) -> None:
  10. self.buffer = bytearray()
  11. self.eof = False
  12. def read_line(self) -> Generator[None, None, bytes]:
  13. """
  14. Read a LF-terminated line from the stream.
  15. The return value includes the LF character.
  16. This is a generator-based coroutine.
  17. :raises EOFError: if the stream ends without a LF
  18. """
  19. n = 0 # number of bytes to read
  20. p = 0 # number of bytes without a newline
  21. while True:
  22. n = self.buffer.find(b"\n", p) + 1
  23. if n > 0:
  24. break
  25. p = len(self.buffer)
  26. if self.eof:
  27. raise EOFError(f"stream ends after {p} bytes, before end of line")
  28. yield
  29. r = self.buffer[:n]
  30. del self.buffer[:n]
  31. return r
  32. def read_exact(self, n: int) -> Generator[None, None, bytes]:
  33. """
  34. Read ``n`` bytes from the stream.
  35. This is a generator-based coroutine.
  36. :raises EOFError: if the stream ends in less than ``n`` bytes
  37. """
  38. assert n >= 0
  39. while len(self.buffer) < n:
  40. if self.eof:
  41. p = len(self.buffer)
  42. raise EOFError(f"stream ends after {p} bytes, expected {n} bytes")
  43. yield
  44. r = self.buffer[:n]
  45. del self.buffer[:n]
  46. return r
  47. def read_to_eof(self) -> Generator[None, None, bytes]:
  48. """
  49. Read all bytes from the stream.
  50. This is a generator-based coroutine.
  51. """
  52. while not self.eof:
  53. yield
  54. r = self.buffer[:]
  55. del self.buffer[:]
  56. return r
  57. def at_eof(self) -> Generator[None, None, bool]:
  58. """
  59. Tell whether the stream has ended and all data was read.
  60. This is a generator-based coroutine.
  61. """
  62. while True:
  63. if self.buffer:
  64. return False
  65. if self.eof:
  66. return True
  67. # When all data was read but the stream hasn't ended, we can't
  68. # tell if until either feed_data() or feed_eof() is called.
  69. yield
  70. def feed_data(self, data: bytes) -> None:
  71. """
  72. Write ``data`` to the stream.
  73. :meth:`feed_data()` cannot be called after :meth:`feed_eof()`.
  74. :raises EOFError: if the stream has ended
  75. """
  76. if self.eof:
  77. raise EOFError("stream ended")
  78. self.buffer += data
  79. def feed_eof(self) -> None:
  80. """
  81. End the stream.
  82. :meth:`feed_eof()` must be called at must once.
  83. :raises EOFError: if the stream has ended
  84. """
  85. if self.eof:
  86. raise EOFError("stream ended")
  87. self.eof = True