From 2c9838e30f825b3bdaf0fc33f15a3cc74320ba6e Mon Sep 17 00:00:00 2001
From: Georg Brandl <georg@python.org>
Date: Sun, 29 Oct 2006 14:39:09 +0000
Subject: [PATCH] Bug #1586613: fix zlib and bz2 codecs' incremental
 en/decoders.

---
 Lib/encodings/bz2_codec.py  | 29 +++++++++++++++++++++++++----
 Lib/encodings/zlib_codec.py | 30 ++++++++++++++++++++++++++----
 Lib/test/test_codecs.py     |  2 ++
 Misc/NEWS                   |  2 ++
 4 files changed, 55 insertions(+), 8 deletions(-)

diff --git a/Lib/encodings/bz2_codec.py b/Lib/encodings/bz2_codec.py
index 81e84b6fb89..054b36b401a 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
diff --git a/Lib/encodings/zlib_codec.py b/Lib/encodings/zlib_codec.py
index 2694f155116..3419f9f48f5 100644
--- a/Lib/encodings/zlib_codec.py
+++ b/Lib/encodings/zlib_codec.py
@@ -51,14 +51,36 @@ class Codec(codecs.Codec):
         return zlib_decode(input, errors)
 
 class IncrementalEncoder(codecs.IncrementalEncoder):
+    def __init__(self, errors='strict'):
+        assert errors == 'strict'
+        self.errors = errors
+        self.compressobj = zlib.compressobj()
+
     def encode(self, input, final=False):
-        assert self.errors == 'strict'
-        return zlib.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 = zlib.compressobj()
 
 class IncrementalDecoder(codecs.IncrementalDecoder):
+    def __init__(self, errors='strict'):
+        assert errors == 'strict'
+        self.errors = errors
+        self.decompressobj = zlib.decompressobj()
+
     def decode(self, input, final=False):
-        assert self.errors == 'strict'
-        return zlib.decompress(input)
+        if final:
+            c = self.decompressobj.decompress(input)
+            return c + self.decompressobj.flush()
+        else:
+            return self.decompressobj.decompress(input)
+
+    def reset(self):
+        self.decompressobj = zlib.decompressobj()
 
 class StreamWriter(Codec,codecs.StreamWriter):
     pass
diff --git a/Lib/test/test_codecs.py b/Lib/test/test_codecs.py
index fa52033ae64..5c82d3f8be6 100644
--- a/Lib/test/test_codecs.py
+++ b/Lib/test/test_codecs.py
@@ -1062,6 +1062,7 @@ broken_unicode_with_streams = [
     "punycode",
     "unicode_internal"
 ]
+broken_incremental_coders = broken_unicode_with_streams[:]
 
 try:
     import bz2
@@ -1111,6 +1112,7 @@ class BasicUnicodeTest(unittest.TestCase):
                     decodedresult += reader.read()
                 self.assertEqual(decodedresult, s, "%r != %r (encoding=%r)" % (decodedresult, s, encoding))
 
+            if encoding not in broken_incremental_coders:
                 # check incremental decoder/encoder (fetched via the Python
                 # and C API) and iterencode()/iterdecode()
                 try:
diff --git a/Misc/NEWS b/Misc/NEWS
index 8363ee37c86..d7548499ff9 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -89,6 +89,8 @@ Core and builtins
 Library
 -------
 
+- Bug #1586613: fix zlib and bz2 codecs' incremental en/decoders.
+
 - Patch #1583880: fix tarfile's problems with long names and posix/
   GNU modes.