jpayne@69: """Python 'bz2_codec' Codec - bz2 compression encoding.
jpayne@69: 
jpayne@69: This codec de/encodes from bytes to bytes and is therefore usable with
jpayne@69: bytes.transform() and bytes.untransform().
jpayne@69: 
jpayne@69: Adapted by Raymond Hettinger from zlib_codec.py which was written
jpayne@69: by Marc-Andre Lemburg (mal@lemburg.com).
jpayne@69: """
jpayne@69: 
jpayne@69: import codecs
jpayne@69: import bz2 # this codec needs the optional bz2 module !
jpayne@69: 
jpayne@69: ### Codec APIs
jpayne@69: 
jpayne@69: def bz2_encode(input, errors='strict'):
jpayne@69:     assert errors == 'strict'
jpayne@69:     return (bz2.compress(input), len(input))
jpayne@69: 
jpayne@69: def bz2_decode(input, errors='strict'):
jpayne@69:     assert errors == 'strict'
jpayne@69:     return (bz2.decompress(input), len(input))
jpayne@69: 
jpayne@69: class Codec(codecs.Codec):
jpayne@69:     def encode(self, input, errors='strict'):
jpayne@69:         return bz2_encode(input, errors)
jpayne@69:     def decode(self, input, errors='strict'):
jpayne@69:         return bz2_decode(input, errors)
jpayne@69: 
jpayne@69: class IncrementalEncoder(codecs.IncrementalEncoder):
jpayne@69:     def __init__(self, errors='strict'):
jpayne@69:         assert errors == 'strict'
jpayne@69:         self.errors = errors
jpayne@69:         self.compressobj = bz2.BZ2Compressor()
jpayne@69: 
jpayne@69:     def encode(self, input, final=False):
jpayne@69:         if final:
jpayne@69:             c = self.compressobj.compress(input)
jpayne@69:             return c + self.compressobj.flush()
jpayne@69:         else:
jpayne@69:             return self.compressobj.compress(input)
jpayne@69: 
jpayne@69:     def reset(self):
jpayne@69:         self.compressobj = bz2.BZ2Compressor()
jpayne@69: 
jpayne@69: class IncrementalDecoder(codecs.IncrementalDecoder):
jpayne@69:     def __init__(self, errors='strict'):
jpayne@69:         assert errors == 'strict'
jpayne@69:         self.errors = errors
jpayne@69:         self.decompressobj = bz2.BZ2Decompressor()
jpayne@69: 
jpayne@69:     def decode(self, input, final=False):
jpayne@69:         try:
jpayne@69:             return self.decompressobj.decompress(input)
jpayne@69:         except EOFError:
jpayne@69:             return ''
jpayne@69: 
jpayne@69:     def reset(self):
jpayne@69:         self.decompressobj = bz2.BZ2Decompressor()
jpayne@69: 
jpayne@69: class StreamWriter(Codec, codecs.StreamWriter):
jpayne@69:     charbuffertype = bytes
jpayne@69: 
jpayne@69: class StreamReader(Codec, codecs.StreamReader):
jpayne@69:     charbuffertype = bytes
jpayne@69: 
jpayne@69: ### encodings module API
jpayne@69: 
jpayne@69: def getregentry():
jpayne@69:     return codecs.CodecInfo(
jpayne@69:         name="bz2",
jpayne@69:         encode=bz2_encode,
jpayne@69:         decode=bz2_decode,
jpayne@69:         incrementalencoder=IncrementalEncoder,
jpayne@69:         incrementaldecoder=IncrementalDecoder,
jpayne@69:         streamwriter=StreamWriter,
jpayne@69:         streamreader=StreamReader,
jpayne@69:         _is_text_encoding=False,
jpayne@69:     )