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