Bug #1586613: fix zlib and bz2 codecs' incremental en/decoders.

 (backport from rev. 52529)

1 files changed, 25 insertions, 4 deletions

diff --git a/Lib/encodings/bz2_codec.py b/Lib/encodings/bz2_codec.py
index 81e84b6fb8..054b36b401 100644
--- a/Lib/encodings/bz2_codec.py
+++ b/Lib/encodings/bz2_codec.py
@@ -52,14 +52,35 @@ class Codec(codecs.Codec):
         return bz2_decode(input, errors)
 
 class IncrementalEncoder(codecs.IncrementalEncoder):
+    def __init__(self, errors='strict'):
+        assert errors == 'strict'
+        self.errors = errors
+        self.compressobj = bz2.BZ2Compressor()
+
     def encode(self, input, final=False):
-        assert self.errors == 'strict'
-        return bz2.compress(input)
+        if final:
+            c = self.compressobj.compress(input)
+            return c + self.compressobj.flush()
+        else:
+            return self.compressobj.compress(input)
+
+    def reset(self):
+        self.compressobj = bz2.BZ2Compressor()
 
 class IncrementalDecoder(codecs.IncrementalDecoder):
+    def __init__(self, errors='strict'):
+        assert errors == 'strict'
+        self.errors = errors
+        self.decompressobj = bz2.BZ2Decompressor()
+
     def decode(self, input, final=False):
-        assert self.errors == 'strict'
-        return bz2.decompress(input)
+        try:
+            return self.decompressobj.decompress(input)
+        except EOFError:
+            return ''
+
+    def reset(self):
+        self.decompressobj = bz2.BZ2Decompressor()
 
 class StreamWriter(Codec,codecs.StreamWriter):
     pass