/lib/pychess/System/readuntil.py

https://github.com/pychess/pychess · Python · 146 lines · 54 code · 22 blank · 70 comment · 12 complexity · 9789698cddbddfe9862d3f7519c24984 MD5 · raw file

  1. """ Monkey patching asyncio.StreamReader to add readuntil() from Python 3.5.2"""
  2. import asyncio
  3. class IncompleteReadError(EOFError):
  4. """
  5. Incomplete read error. Attributes:
  6. - partial: read bytes string before the end of stream was reached
  7. - expected: total number of expected bytes (or None if unknown)
  8. """
  9. def __init__(self, partial, expected):
  10. super().__init__("%d bytes read on a total of %r expected bytes"
  11. % (len(partial), expected))
  12. self.partial = partial
  13. self.expected = expected
  14. class LimitOverrunError(Exception):
  15. """Reached the buffer limit while looking for a separator.
  16. Attributes:
  17. - consumed: total number of to be consumed bytes.
  18. """
  19. def __init__(self, message, consumed):
  20. super().__init__(message)
  21. self.consumed = consumed
  22. async def _wait_for_data(self, func_name):
  23. """Wait until feed_data() or feed_eof() is called.
  24. If stream was paused, automatically resume it.
  25. """
  26. # StreamReader uses a future to link the protocol feed_data() method
  27. # to a read coroutine. Running two read coroutines at the same time
  28. # would have an unexpected behaviour. It would not possible to know
  29. # which coroutine would get the next data.
  30. if self._waiter is not None:
  31. raise RuntimeError('%s() called while another coroutine is '
  32. 'already waiting for incoming data' % func_name)
  33. assert not self._eof, '_wait_for_data after EOF'
  34. # Waiting for data while paused will make deadlock, so prevent it.
  35. if self._paused:
  36. self._paused = False
  37. self._transport.resume_reading()
  38. self._waiter = asyncio.futures.Future(loop=self._loop)
  39. try:
  40. await self._waiter
  41. finally:
  42. self._waiter = None
  43. async def readuntil(self, separator=b'\n'):
  44. """Read data from the stream until ``separator`` is found.
  45. On success, the data and separator will be removed from the
  46. internal buffer (consumed). Returned data will include the
  47. separator at the end.
  48. Configured stream limit is used to check result. Limit sets the
  49. maximal length of data that can be returned, not counting the
  50. separator.
  51. If an EOF occurs and the complete separator is still not found,
  52. an IncompleteReadError exception will be raised, and the internal
  53. buffer will be reset. The IncompleteReadError.partial attribute
  54. may contain the separator partially.
  55. If the data cannot be read because of over limit, a
  56. LimitOverrunError exception will be raised, and the data
  57. will be left in the internal buffer, so it can be read again.
  58. """
  59. seplen = len(separator)
  60. if seplen == 0:
  61. raise ValueError('Separator should be at least one-byte string')
  62. if self._exception is not None:
  63. raise self._exception
  64. # Consume whole buffer except last bytes, which length is
  65. # one less than seplen. Let's check corner cases with
  66. # separator='SEPARATOR':
  67. # * we have received almost complete separator (without last
  68. # byte). i.e buffer='some textSEPARATO'. In this case we
  69. # can safely consume len(separator) - 1 bytes.
  70. # * last byte of buffer is first byte of separator, i.e.
  71. # buffer='abcdefghijklmnopqrS'. We may safely consume
  72. # everything except that last byte, but this require to
  73. # analyze bytes of buffer that match partial separator.
  74. # This is slow and/or require FSM. For this case our
  75. # implementation is not optimal, since require rescanning
  76. # of data that is known to not belong to separator. In
  77. # real world, separator will not be so long to notice
  78. # performance problems. Even when reading MIME-encoded
  79. # messages :)
  80. # `offset` is the number of bytes from the beginning of the buffer
  81. # where there is no occurrence of `separator`.
  82. offset = 0
  83. # Loop until we find `separator` in the buffer, exceed the buffer size,
  84. # or an EOF has happened.
  85. while True:
  86. buflen = len(self._buffer)
  87. # Check if we now have enough data in the buffer for `separator` to
  88. # fit.
  89. if buflen - offset >= seplen:
  90. isep = self._buffer.find(separator, offset)
  91. if isep != -1:
  92. # `separator` is in the buffer. `isep` will be used later
  93. # to retrieve the data.
  94. break
  95. # see upper comment for explanation.
  96. offset = buflen + 1 - seplen
  97. if offset > self._limit:
  98. raise LimitOverrunError(
  99. 'Separator is not found, and chunk exceed the limit',
  100. offset)
  101. # Complete message (with full separator) may be present in buffer
  102. # even when EOF flag is set. This may happen when the last chunk
  103. # adds data which makes separator be found. That's why we check for
  104. # EOF *ater* inspecting the buffer.
  105. if self._eof:
  106. chunk = bytes(self._buffer)
  107. self._buffer.clear()
  108. raise IncompleteReadError(chunk, None)
  109. # _wait_for_data() will resume reading if stream was paused.
  110. await self._wait_for_data('readuntil')
  111. if isep > self._limit:
  112. raise LimitOverrunError(
  113. 'Separator is found, but chunk is longer than limit', isep)
  114. chunk = self._buffer[:isep + seplen]
  115. del self._buffer[:isep + seplen]
  116. self._maybe_resume_transport()
  117. return bytes(chunk)