_decoders.py 9.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327
  1. """
  2. Handlers for Content-Encoding.
  3. See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Encoding
  4. """
  5. import codecs
  6. import io
  7. import typing
  8. import zlib
  9. from ._compat import brotli
  10. from ._exceptions import DecodingError
  11. class ContentDecoder:
  12. def decode(self, data: bytes) -> bytes:
  13. raise NotImplementedError() # pragma: no cover
  14. def flush(self) -> bytes:
  15. raise NotImplementedError() # pragma: no cover
  16. class IdentityDecoder(ContentDecoder):
  17. """
  18. Handle unencoded data.
  19. """
  20. def decode(self, data: bytes) -> bytes:
  21. return data
  22. def flush(self) -> bytes:
  23. return b""
  24. class DeflateDecoder(ContentDecoder):
  25. """
  26. Handle 'deflate' decoding.
  27. See: https://stackoverflow.com/questions/1838699
  28. """
  29. def __init__(self) -> None:
  30. self.first_attempt = True
  31. self.decompressor = zlib.decompressobj()
  32. def decode(self, data: bytes) -> bytes:
  33. was_first_attempt = self.first_attempt
  34. self.first_attempt = False
  35. try:
  36. return self.decompressor.decompress(data)
  37. except zlib.error as exc:
  38. if was_first_attempt:
  39. self.decompressor = zlib.decompressobj(-zlib.MAX_WBITS)
  40. return self.decode(data)
  41. raise DecodingError(str(exc)) from exc
  42. def flush(self) -> bytes:
  43. try:
  44. return self.decompressor.flush()
  45. except zlib.error as exc: # pragma: no cover
  46. raise DecodingError(str(exc)) from exc
  47. class GZipDecoder(ContentDecoder):
  48. """
  49. Handle 'gzip' decoding.
  50. See: https://stackoverflow.com/questions/1838699
  51. """
  52. def __init__(self) -> None:
  53. self.decompressor = zlib.decompressobj(zlib.MAX_WBITS | 16)
  54. def decode(self, data: bytes) -> bytes:
  55. try:
  56. return self.decompressor.decompress(data)
  57. except zlib.error as exc:
  58. raise DecodingError(str(exc)) from exc
  59. def flush(self) -> bytes:
  60. try:
  61. return self.decompressor.flush()
  62. except zlib.error as exc: # pragma: no cover
  63. raise DecodingError(str(exc)) from exc
  64. class BrotliDecoder(ContentDecoder):
  65. """
  66. Handle 'brotli' decoding.
  67. Requires `pip install brotlipy`. See: https://brotlipy.readthedocs.io/
  68. or `pip install brotli`. See https://github.com/google/brotli
  69. Supports both 'brotlipy' and 'Brotli' packages since they share an import
  70. name. The top branches are for 'brotlipy' and bottom branches for 'Brotli'
  71. """
  72. def __init__(self) -> None:
  73. if brotli is None: # pragma: no cover
  74. raise ImportError(
  75. "Using 'BrotliDecoder', but neither of the 'brotlicffi' or 'brotli' "
  76. "packages have been installed. "
  77. "Make sure to install httpx using `pip install httpx[brotli]`."
  78. ) from None
  79. self.decompressor = brotli.Decompressor()
  80. self.seen_data = False
  81. self._decompress: typing.Callable[[bytes], bytes]
  82. if hasattr(self.decompressor, "decompress"):
  83. # The 'brotlicffi' package.
  84. self._decompress = self.decompressor.decompress # pragma: no cover
  85. else:
  86. # The 'brotli' package.
  87. self._decompress = self.decompressor.process # pragma: no cover
  88. def decode(self, data: bytes) -> bytes:
  89. if not data:
  90. return b""
  91. self.seen_data = True
  92. try:
  93. return self._decompress(data)
  94. except brotli.error as exc:
  95. raise DecodingError(str(exc)) from exc
  96. def flush(self) -> bytes:
  97. if not self.seen_data:
  98. return b""
  99. try:
  100. if hasattr(self.decompressor, "finish"):
  101. # Only available in the 'brotlicffi' package.
  102. # As the decompressor decompresses eagerly, this
  103. # will never actually emit any data. However, it will potentially throw
  104. # errors if a truncated or damaged data stream has been used.
  105. self.decompressor.finish() # pragma: no cover
  106. return b""
  107. except brotli.error as exc: # pragma: no cover
  108. raise DecodingError(str(exc)) from exc
  109. class MultiDecoder(ContentDecoder):
  110. """
  111. Handle the case where multiple encodings have been applied.
  112. """
  113. def __init__(self, children: typing.Sequence[ContentDecoder]) -> None:
  114. """
  115. 'children' should be a sequence of decoders in the order in which
  116. each was applied.
  117. """
  118. # Note that we reverse the order for decoding.
  119. self.children = list(reversed(children))
  120. def decode(self, data: bytes) -> bytes:
  121. for child in self.children:
  122. data = child.decode(data)
  123. return data
  124. def flush(self) -> bytes:
  125. data = b""
  126. for child in self.children:
  127. data = child.decode(data) + child.flush()
  128. return data
  129. class ByteChunker:
  130. """
  131. Handles returning byte content in fixed-size chunks.
  132. """
  133. def __init__(self, chunk_size: typing.Optional[int] = None) -> None:
  134. self._buffer = io.BytesIO()
  135. self._chunk_size = chunk_size
  136. def decode(self, content: bytes) -> typing.List[bytes]:
  137. if self._chunk_size is None:
  138. return [content] if content else []
  139. self._buffer.write(content)
  140. if self._buffer.tell() >= self._chunk_size:
  141. value = self._buffer.getvalue()
  142. chunks = [
  143. value[i : i + self._chunk_size]
  144. for i in range(0, len(value), self._chunk_size)
  145. ]
  146. if len(chunks[-1]) == self._chunk_size:
  147. self._buffer.seek(0)
  148. self._buffer.truncate()
  149. return chunks
  150. else:
  151. self._buffer.seek(0)
  152. self._buffer.write(chunks[-1])
  153. self._buffer.truncate()
  154. return chunks[:-1]
  155. else:
  156. return []
  157. def flush(self) -> typing.List[bytes]:
  158. value = self._buffer.getvalue()
  159. self._buffer.seek(0)
  160. self._buffer.truncate()
  161. return [value] if value else []
  162. class TextChunker:
  163. """
  164. Handles returning text content in fixed-size chunks.
  165. """
  166. def __init__(self, chunk_size: typing.Optional[int] = None) -> None:
  167. self._buffer = io.StringIO()
  168. self._chunk_size = chunk_size
  169. def decode(self, content: str) -> typing.List[str]:
  170. if self._chunk_size is None:
  171. return [content] if content else []
  172. self._buffer.write(content)
  173. if self._buffer.tell() >= self._chunk_size:
  174. value = self._buffer.getvalue()
  175. chunks = [
  176. value[i : i + self._chunk_size]
  177. for i in range(0, len(value), self._chunk_size)
  178. ]
  179. if len(chunks[-1]) == self._chunk_size:
  180. self._buffer.seek(0)
  181. self._buffer.truncate()
  182. return chunks
  183. else:
  184. self._buffer.seek(0)
  185. self._buffer.write(chunks[-1])
  186. self._buffer.truncate()
  187. return chunks[:-1]
  188. else:
  189. return []
  190. def flush(self) -> typing.List[str]:
  191. value = self._buffer.getvalue()
  192. self._buffer.seek(0)
  193. self._buffer.truncate()
  194. return [value] if value else []
  195. class TextDecoder:
  196. """
  197. Handles incrementally decoding bytes into text
  198. """
  199. def __init__(self, encoding: str = "utf-8") -> None:
  200. self.decoder = codecs.getincrementaldecoder(encoding)(errors="replace")
  201. def decode(self, data: bytes) -> str:
  202. return self.decoder.decode(data)
  203. def flush(self) -> str:
  204. return self.decoder.decode(b"", True)
  205. class LineDecoder:
  206. """
  207. Handles incrementally reading lines from text.
  208. Has the same behaviour as the stdllib splitlines,
  209. but handling the input iteratively.
  210. """
  211. def __init__(self) -> None:
  212. self.buffer: typing.List[str] = []
  213. self.trailing_cr: bool = False
  214. def decode(self, text: str) -> typing.List[str]:
  215. # See https://docs.python.org/3/library/stdtypes.html#str.splitlines
  216. NEWLINE_CHARS = "\n\r\x0b\x0c\x1c\x1d\x1e\x85\u2028\u2029"
  217. # We always push a trailing `\r` into the next decode iteration.
  218. if self.trailing_cr:
  219. text = "\r" + text
  220. self.trailing_cr = False
  221. if text.endswith("\r"):
  222. self.trailing_cr = True
  223. text = text[:-1]
  224. if not text:
  225. # NOTE: the edge case input of empty text doesn't occur in practice,
  226. # because other httpx internals filter out this value
  227. return [] # pragma: no cover
  228. trailing_newline = text[-1] in NEWLINE_CHARS
  229. lines = text.splitlines()
  230. if len(lines) == 1 and not trailing_newline:
  231. # No new lines, buffer the input and continue.
  232. self.buffer.append(lines[0])
  233. return []
  234. if self.buffer:
  235. # Include any existing buffer in the first portion of the
  236. # splitlines result.
  237. lines = ["".join(self.buffer) + lines[0]] + lines[1:]
  238. self.buffer = []
  239. if not trailing_newline:
  240. # If the last segment of splitlines is not newline terminated,
  241. # then drop it from our output and start a new buffer.
  242. self.buffer = [lines.pop()]
  243. return lines
  244. def flush(self) -> typing.List[str]:
  245. if not self.buffer and not self.trailing_cr:
  246. return []
  247. lines = ["".join(self.buffer)]
  248. self.buffer = []
  249. self.trailing_cr = False
  250. return lines
  251. SUPPORTED_DECODERS = {
  252. "identity": IdentityDecoder,
  253. "gzip": GZipDecoder,
  254. "deflate": DeflateDecoder,
  255. "br": BrotliDecoder,
  256. }
  257. if brotli is None:
  258. SUPPORTED_DECODERS.pop("br") # pragma: no cover