codekingpro/portable-devtools

1"""2Utility functions for decoding response bodies.3"""4 5import codecs6import collections7import gzip8import zlib9from io import BytesIO10from typing import overload11 12import brotli13import zstandard as zstd14 15# We have a shared single-element cache for encoding and decoding.16# This is quite useful in practice, e.g.17# flow.request.content = flow.request.content.replace(b"foo", b"bar")18# does not require an .encode() call if content does not contain b"foo"19CachedDecode = collections.namedtuple("CachedDecode", "encoded encoding errors decoded")20_cache = CachedDecode(None, None, None, None)21 22 23@overload24def decode(encoded: None, encoding: str, errors: str = "strict") -> None: ...25 26 27@overload28def decode(encoded: str, encoding: str, errors: str = "strict") -> str: ...29 30 31@overload32def decode(encoded: bytes, encoding: str, errors: str = "strict") -> str | bytes: ...33 34 35def decode(36    encoded: None | str | bytes, encoding: str, errors: str = "strict"37) -> None | str | bytes:38    """39    Decode the given input object40 41    Returns:42        The decoded value43 44    Raises:45        ValueError, if decoding fails.46    """47    if encoded is None:48        return None49    encoding = encoding.lower()50 51    global _cache52    cached = (53        isinstance(encoded, bytes)54        and _cache.encoded == encoded55        and _cache.encoding == encoding56        and _cache.errors == errors57    )58    if cached:59        return _cache.decoded60    try:61        try:62            decoded = custom_decode[encoding](encoded)63        except KeyError:64            decoded = codecs.decode(encoded, encoding, errors)  # type: ignore65        if encoding in ("gzip", "deflate", "deflateraw", "br", "zstd"):66            _cache = CachedDecode(encoded, encoding, errors, decoded)67        return decoded68    except TypeError:69        raise70    except Exception as e:71        raise ValueError(72            "{} when decoding {} with {}: {}".format(73                type(e).__name__,74                repr(encoded)[:10],75                repr(encoding),76                repr(e),77            )78        )79 80 81@overload82def encode(decoded: None, encoding: str, errors: str = "strict") -> None: ...83 84 85@overload86def encode(decoded: str, encoding: str, errors: str = "strict") -> str | bytes: ...87 88 89@overload90def encode(decoded: bytes, encoding: str, errors: str = "strict") -> bytes: ...91 92 93def encode(94    decoded: None | str | bytes, encoding, errors="strict"95) -> None | str | bytes:96    """97    Encode the given input object98 99    Returns:100        The encoded value101 102    Raises:103        ValueError, if encoding fails.104    """105    if decoded is None:106        return None107    encoding = encoding.lower()108 109    global _cache110    cached = (111        isinstance(decoded, bytes)112        and _cache.decoded == decoded113        and _cache.encoding == encoding114        and _cache.errors == errors115    )116    if cached:117        return _cache.encoded118    try:119        try:120            encoded = custom_encode[encoding](decoded)121        except KeyError:122            encoded = codecs.encode(decoded, encoding, errors)  # type: ignore123        if encoding in ("gzip", "deflate", "deflateraw", "br", "zstd"):124            _cache = CachedDecode(encoded, encoding, errors, decoded)125        return encoded126    except TypeError:127        raise128    except Exception as e:129        raise ValueError(130            "{} when encoding {} with {}: {}".format(131                type(e).__name__,132                repr(decoded)[:10],133                repr(encoding),134                repr(e),135            )136        )137 138 139def identity(content):140    """141    Returns content unchanged. Identity is the default value of142    Accept-Encoding headers.143    """144    return content145 146 147def decode_gzip(content: bytes) -> bytes:148    """Decode gzip or zlib-compressed data using zlib's auto-detection."""149    if not content:150        return b""151 152    try:153        # Using wbits=47 (32 + 15) tells zlib to automatically detect both gzip and zlib headers.154        # This simplifies decoding and avoids the need for a separate gzip.GzipFile fallback.155        # Reference: https://docs.python.org/3/library/zlib.html#zlib.decompress156        decompressor = zlib.decompressobj(47)157        return decompressor.decompress(content) + decompressor.flush()158    except zlib.error as e:159        raise ValueError(f"Decompression failed: {e}")160 161 162def encode_gzip(content: bytes) -> bytes:163    s = BytesIO()164    # set mtime to 0 so that gzip encoding is deterministic.165    # Use compresslevel=1 for fastest compression speed.166    with gzip.GzipFile(fileobj=s, mode="wb", mtime=0, compresslevel=1) as f:167        f.write(content)168    return s.getvalue()169 170 171def decode_brotli(content: bytes) -> bytes:172    if not content:173        return b""174    return brotli.decompress(content)175 176 177def encode_brotli(content: bytes) -> bytes:178    # Use quality=0 for fastest compression speed.179    return brotli.compress(content, quality=0)180 181 182def decode_zstd(content: bytes) -> bytes:183    if not content:184        return b""185    zstd_ctx = zstd.ZstdDecompressor()186    return zstd_ctx.stream_reader(BytesIO(content), read_across_frames=True).read()187 188 189def encode_zstd(content: bytes) -> bytes:190    # Use level=1 for fastest compression speed.191    zstd_ctx = zstd.ZstdCompressor(level=1)192    return zstd_ctx.compress(content)193 194 195def decode_deflate(content: bytes) -> bytes:196    """197    Returns decompressed data for DEFLATE. Some servers may respond with198    compressed data without a zlib header or checksum. An undocumented199    feature of zlib permits the lenient decompression of data missing both200    values.201 202    http://bugs.python.org/issue5784203    """204    if not content:205        return b""206    try:207        return zlib.decompress(content)208    except zlib.error:209        return zlib.decompress(content, -15)210 211 212def encode_deflate(content: bytes) -> bytes:213    """214    Returns compressed content, always including zlib header and checksum.215    """216    # Use level=1 for fastest compression speed.217    return zlib.compress(content, level=1)218 219 220custom_decode = {221    "none": identity,222    "identity": identity,223    "gzip": decode_gzip,224    "deflate": decode_deflate,225    "deflateraw": decode_deflate,226    "br": decode_brotli,227    "zstd": decode_zstd,228}229custom_encode = {230    "none": identity,231    "identity": identity,232    "gzip": encode_gzip,233    "deflate": encode_deflate,234    "deflateraw": encode_deflate,235    "br": encode_brotli,236    "zstd": encode_zstd,237}238 239__all__ = ["encode", "decode"]240