| 248 | HAS_ZSTD = True |
| 249 | |
| 250 | class ZstdDecoder(ContentDecoder): |
| 251 | def __init__(self) -> None: |
| 252 | self._obj = zstd.ZstdDecompressor() |
| 253 | |
| 254 | def decompress(self, data: bytes, max_length: int = -1) -> bytes: |
| 255 | if not data and not self.has_unconsumed_tail: |
| 256 | return b"" |
| 257 | if self._obj.eof: |
| 258 | data = self._obj.unused_data + data |
| 259 | self._obj = zstd.ZstdDecompressor() |
| 260 | part = self._obj.decompress(data, max_length=max_length) |
| 261 | length = len(part) |
| 262 | data_parts = [part] |
| 263 | # Every loop iteration is supposed to read data from a separate frame. |
| 264 | # The loop breaks when: |
| 265 | # - enough data is read; |
| 266 | # - no more unused data is available; |
| 267 | # - end of the last read frame has not been reached (i.e., |
| 268 | # more data has to be fed). |
| 269 | while ( |
| 270 | self._obj.eof |
| 271 | and self._obj.unused_data |
| 272 | and (max_length < 0 or length < max_length) |
| 273 | ): |
| 274 | unused_data = self._obj.unused_data |
| 275 | if not self._obj.needs_input: |
| 276 | self._obj = zstd.ZstdDecompressor() |
| 277 | part = self._obj.decompress( |
| 278 | unused_data, |
| 279 | max_length=(max_length - length) if max_length > 0 else -1, |
| 280 | ) |
| 281 | if part_length := len(part): |
| 282 | data_parts.append(part) |
| 283 | length += part_length |
| 284 | elif self._obj.needs_input: |
| 285 | break |
| 286 | return b"".join(data_parts) |
| 287 | |
| 288 | @property |
| 289 | def has_unconsumed_tail(self) -> bool: |
| 290 | return not (self._obj.needs_input or self._obj.eof) or bool( |
| 291 | self._obj.unused_data |
| 292 | ) |
| 293 | |
| 294 | def flush(self) -> bytes: |
| 295 | if not self._obj.eof: |
| 296 | raise DecodeError("Zstandard data is incomplete") |
| 297 | return b"" |
| 298 | |
| 299 | |
| 300 | class MultiDecoder(ContentDecoder): |