1"""Python 'zlib_codec' Codec - zlib compression encoding.
2
3This codec de/encodes from bytes to bytes.
4
5Written by Marc-Andre Lemburg (mal@lemburg.com).
6"""
7
8import codecs
9import zlib # this codec needs the optional zlib module !
10
11### Codec APIs
12
13def zlib_encode(input, errors='strict'):
14    assert errors == 'strict'
15    return (zlib.compress(input), len(input))
16
17def zlib_decode(input, errors='strict'):
18    assert errors == 'strict'
19    return (zlib.decompress(input), len(input))
20
21class Codec(codecs.Codec):
22    def encode(self, input, errors='strict'):
23        return zlib_encode(input, errors)
24    def decode(self, input, errors='strict'):
25        return zlib_decode(input, errors)
26
27class IncrementalEncoder(codecs.IncrementalEncoder):
28    def __init__(self, errors='strict'):
29        assert errors == 'strict'
30        self.errors = errors
31        self.compressobj = zlib.compressobj()
32
33    def encode(self, input, final=False):
34        if final:
35            c = self.compressobj.compress(input)
36            return c + self.compressobj.flush()
37        else:
38            return self.compressobj.compress(input)
39
40    def reset(self):
41        self.compressobj = zlib.compressobj()
42
43class IncrementalDecoder(codecs.IncrementalDecoder):
44    def __init__(self, errors='strict'):
45        assert errors == 'strict'
46        self.errors = errors
47        self.decompressobj = zlib.decompressobj()
48
49    def decode(self, input, final=False):
50        if final:
51            c = self.decompressobj.decompress(input)
52            return c + self.decompressobj.flush()
53        else:
54            return self.decompressobj.decompress(input)
55
56    def reset(self):
57        self.decompressobj = zlib.decompressobj()
58
59class StreamWriter(Codec, codecs.StreamWriter):
60    charbuffertype = bytes
61
62class StreamReader(Codec, codecs.StreamReader):
63    charbuffertype = bytes
64
65### encodings module API
66
67def getregentry():
68    return codecs.CodecInfo(
69        name='zlib',
70        encode=zlib_encode,
71        decode=zlib_decode,
72        incrementalencoder=IncrementalEncoder,
73        incrementaldecoder=IncrementalDecoder,
74        streamreader=StreamReader,
75        streamwriter=StreamWriter,
76        _is_text_encoding=False,
77    )
78