Issue #5784: Expand documentation and tests for zlib wbits parameter

Based on documentation by AM Kuchling.

3 files changed, 98 insertions, 19 deletions

diff --git a/Doc/library/zlib.rst b/Doc/library/zlib.rst
index 26b0dfc..1e2df00 100644
--- a/Doc/library/zlib.rst
+++ b/Doc/library/zlib.rst
@@ -81,9 +81,23 @@ The available exception and functions in this module are:
    *method* is the compression algorithm. Currently, the only supported value is
    ``DEFLATED``.
 
-   *wbits* is the base two logarithm of the size of the window buffer. This
-   should be an integer from ``8`` to ``15``. Higher values give better
-   compression, but use more memory. The default is 15.
+   The *wbits* argument controls the size of the history buffer (or the
+   "window size") used when compressing data, and whether a header and
+   trailer is included in the output.  It can take several ranges of values.
+   The default is 15.
+
+   * +9 to +15: The base-two logarithm of the window size, which
+     therefore ranges between 512 and 32768.  Larger values produce
+     better compression at the expense of greater memory usage.  The
+     resulting output will include a zlib-specific header and trailer.
+
+   * −9 to −15: Uses the absolute value of *wbits* as the
+     window size logarithm, while producing a raw output stream with no
+     header or trailing checksum.
+
+   * +25 to +31 = 16 + (9 to 15): Uses the low 4 bits of the value as the
+     window size logarithm, while including a basic :program:`gzip` header
+     and trailing checksum in the output.
 
    *memlevel* controls the amount of memory used for internal compression state.
    Valid values range from ``1`` to ``9``. Higher values using more memory,
@@ -130,20 +144,39 @@ The available exception and functions in this module are:
 .. function:: decompress(string[, wbits[, bufsize]])
 
    Decompresses the data in *string*, returning a string containing the
-   uncompressed data.  The *wbits* parameter controls the size of the window
-   buffer, and is discussed further below.
+   uncompressed data.  The *wbits* parameter depends on
+   the format of *string*, and is discussed further below.
    If *bufsize* is given, it is used as the initial size of the output
    buffer.  Raises the :exc:`error` exception if any error occurs.
 
-   The absolute value of *wbits* is the base two logarithm of the size of the
-   history buffer (the "window size") used when compressing data.  Its absolute
-   value should be between 8 and 15 for the most recent versions of the zlib
-   library, larger values resulting in better compression at the expense of greater
-   memory usage.  When decompressing a stream, *wbits* must not be smaller
+   .. _decompress-wbits:
+
+   The *wbits* parameter controls the size of the history buffer
+   (or "window size"), and what header and trailer format is expected.
+   It is similar to the parameter for :func:`compressobj`, but accepts
+   more ranges of values:
+
+   * +8 to +15: The base-two logarithm of the window size.  The input
+     must include a zlib header and trailer.
+
+   * 0: Automatically determine the window size from the zlib header.
+
+   * −8 to −15: Uses the absolute value of *wbits* as the window size
+     logarithm.  The input must be a raw stream with no header or trailer.
+
+   * +24 to +31 = 16 + (8 to 15): Uses the low 4 bits of the value as
+     the window size logarithm.  The input must include a gzip header and
+     trailer.
+
+   * +40 to +47 = 32 + (8 to 15): Uses the low 4 bits of the value as
+     the window size logarithm, and automatically accepts either
+     the zlib or gzip format.
+
+   When decompressing a stream, the window size must not be smaller
    than the size originally used to compress the stream; using a too-small
-   value will result in an exception. The default value is therefore the
-   highest value, 15.  When *wbits* is negative, the standard
-   :program:`gzip` header is suppressed.
+   value may result in an :exc:`error` exception. The default *wbits* value
+   is 15, which corresponds to the largest window size and requires a zlib
+   header and trailer to be included.
 
    *bufsize* is the initial size of the buffer used to hold decompressed data.  If
    more space is required, the buffer size will be increased as needed, so you
@@ -154,8 +187,11 @@ The available exception and functions in this module are:
 .. function:: decompressobj([wbits])
 
    Returns a decompression object, to be used for decompressing data streams that
-   won't fit into memory at once.  The *wbits* parameter controls the size of the
-   window buffer.
+   won't fit into memory at once.
+
+   The *wbits* parameter controls the size of the history buffer (or the
+   "window size"), and what header and trailer format is expected.  It has
+   the same meaning as `described for decompress() <#decompress-wbits>`__.
 
 Compression objects support the following methods:
 
diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
index cc13409..c537c2a 100644
--- a/Lib/test/test_zlib.py
+++ b/Lib/test/test_zlib.py
@@ -527,6 +527,49 @@ class CompressObjectTestCase(BaseCompressTestCase, unittest.TestCase):
         decompress = lambda s: d.decompress(s) + d.flush()
         self.check_big_decompress_buffer(size, decompress)
 
+    def test_wbits(self):
+        co = zlib.compressobj(1, zlib.DEFLATED, 15)
+        zlib15 = co.compress(HAMLET_SCENE) + co.flush()
+        self.assertEqual(zlib.decompress(zlib15, 15), HAMLET_SCENE)
+        self.assertEqual(zlib.decompress(zlib15, 0), HAMLET_SCENE)
+        self.assertEqual(zlib.decompress(zlib15, 32 + 15), HAMLET_SCENE)
+        with self.assertRaisesRegexp(zlib.error, 'invalid window size'):
+            zlib.decompress(zlib15, 14)
+        dco = zlib.decompressobj(32 + 15)
+        self.assertEqual(dco.decompress(zlib15), HAMLET_SCENE)
+        dco = zlib.decompressobj(14)
+        with self.assertRaisesRegexp(zlib.error, 'invalid window size'):
+            dco.decompress(zlib15)
+
+        co = zlib.compressobj(1, zlib.DEFLATED, 9)
+        zlib9 = co.compress(HAMLET_SCENE) + co.flush()
+        self.assertEqual(zlib.decompress(zlib9, 9), HAMLET_SCENE)
+        self.assertEqual(zlib.decompress(zlib9, 15), HAMLET_SCENE)
+        self.assertEqual(zlib.decompress(zlib9, 0), HAMLET_SCENE)
+        self.assertEqual(zlib.decompress(zlib9, 32 + 9), HAMLET_SCENE)
+        dco = zlib.decompressobj(32 + 9)
+        self.assertEqual(dco.decompress(zlib9), HAMLET_SCENE)
+
+        co = zlib.compressobj(1, zlib.DEFLATED, -15)
+        deflate15 = co.compress(HAMLET_SCENE) + co.flush()
+        self.assertEqual(zlib.decompress(deflate15, -15), HAMLET_SCENE)
+        dco = zlib.decompressobj(-15)
+        self.assertEqual(dco.decompress(deflate15), HAMLET_SCENE)
+
+        co = zlib.compressobj(1, zlib.DEFLATED, -9)
+        deflate9 = co.compress(HAMLET_SCENE) + co.flush()
+        self.assertEqual(zlib.decompress(deflate9, -9), HAMLET_SCENE)
+        self.assertEqual(zlib.decompress(deflate9, -15), HAMLET_SCENE)
+        dco = zlib.decompressobj(-9)
+        self.assertEqual(dco.decompress(deflate9), HAMLET_SCENE)
+
+        co = zlib.compressobj(1, zlib.DEFLATED, 16 + 15)
+        gzip = co.compress(HAMLET_SCENE) + co.flush()
+        self.assertEqual(zlib.decompress(gzip, 16 + 15), HAMLET_SCENE)
+        self.assertEqual(zlib.decompress(gzip, 32 + 15), HAMLET_SCENE)
+        dco = zlib.decompressobj(32 + 15)
+        self.assertEqual(dco.decompress(gzip), HAMLET_SCENE)
+
 
 def genblock(seed, length, step=1024, generator=random):
     """length-byte stream of random data from a seed (in step-byte blocks)."""
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
index 799e7df..153a86c 100644
--- a/Modules/zlibmodule.c
+++ b/Modules/zlibmodule.c
@@ -106,7 +106,7 @@ PyDoc_STRVAR(compressobj__doc__,
 PyDoc_STRVAR(decompressobj__doc__,
 "decompressobj([wbits]) -- Return a decompressor object.\n"
 "\n"
-"Optional arg wbits is the window buffer size.");
+"Optional arg wbits indicates the window buffer size and container format.");
 
 static compobject *
 newcompobject(PyTypeObject *type)
@@ -208,8 +208,8 @@ PyZlib_compress(PyObject *self, PyObject *args)
 PyDoc_STRVAR(decompress__doc__,
 "decompress(string[, wbits[, bufsize]]) -- Return decompressed string.\n"
 "\n"
-"Optional arg wbits is the window buffer size.  Optional arg bufsize is\n"
-"the initial output buffer size.");
+"Optional arg wbits indicates the window buffer size and container format.\n"
+"Optional arg bufsize is the initial output buffer size.");
 
 static PyObject *
 PyZlib_decompress(PyObject *self, PyObject *args)
@@ -1039,7 +1039,7 @@ PyDoc_STRVAR(zlib_module_documentation,
 "decompress(string,[wbits],[bufsize]) -- Decompresses a compressed string.\n"
 "decompressobj([wbits]) -- Return a decompressor object.\n"
 "\n"
-"'wbits' is window buffer size.\n"
+"'wbits' is window buffer size and container format.\n"
 "Compressor objects support compress() and flush() methods; decompressor\n"
 "objects support decompress() and flush().");