python · rhpvorderman · Mar 24, 2021 · Mar 24, 2021 · Mar 24, 2021 · Mar 24, 2021
diff --git a/Doc/library/zlib.rst b/Doc/library/zlib.rst
@@ -47,19 +47,43 @@ The available exception and functions in this module are:
      platforms, use ``adler32(data) & 0xffffffff``.


-.. function:: compress(data, /, level=-1)
+.. function:: compress(data, /, level=-1, wbits=MAX_WBITS)

   Compresses the bytes in *data*, returning a bytes object containing compressed data.
   *level* is an integer from ``0`` to ``9`` or ``-1`` controlling the level of compression;
   ``1`` (Z_BEST_SPEED) is fastest and produces the least compression, ``9`` (Z_BEST_COMPRESSION)
   is slowest and produces the most.  ``0`` (Z_NO_COMPRESSION) is no compression.
   The default value is ``-1`` (Z_DEFAULT_COMPRESSION).  Z_DEFAULT_COMPRESSION represents a default
   compromise between speed and compression (currently equivalent to level 6).
+
+   .. _compress-wbits:
+
+   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,
+   defaulting to ``15`` (MAX_WBITS):
+
+   * +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.
+
   Raises the :exc:`error` exception if any error occurs.

   .. versionchanged:: 3.6
      *level* can now be used as a keyword parameter.

+   .. versionchanged:: 3.11
+      The *wbits* parameter is now available to set window bits and
+      compression type.

 .. function:: compressobj(level=-1, method=DEFLATED, wbits=MAX_WBITS, memLevel=DEF_MEM_LEVEL, strategy=Z_DEFAULT_STRATEGY[, zdict])

@@ -76,23 +100,9 @@ The available exception and functions in this module are:
   *method* is the compression algorithm. Currently, the only supported value is
   :const:`DEFLATED`.

-   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,
-   defaulting to ``15`` (MAX_WBITS):
-
-   * +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.
+   The *wbits* parameter controls the size of the history buffer (or the
+   "window size"), and what header and trailer format will be used. It has
+   the same meaning as `described for compress() <#compress-wbits>`__.

   The *memLevel* argument controls the amount of memory used for the
   internal compression state. Valid values range from ``1`` to ``9``.

diff --git a/Lib/test/test_zlib.py b/Lib/test/test_zlib.py
@@ -825,6 +825,12 @@ def test_wbits(self):
        dco = zlib.decompressobj(32 + 15)
        self.assertEqual(dco.decompress(gzip), HAMLET_SCENE)

+        for wbits in (-15, 15, 31):
+            self.assertEqual(
+                zlib.decompress(
+                    zlib.compress(HAMLET_SCENE, wbits=wbits), wbits=wbits
+                ), HAMLET_SCENE)
+

 def choose_lines(source, number, seed=None, generator=random):
    """Return a list of number lines randomly chosen from the source"""

diff --git a/Misc/NEWS.d/next/Library/2021-03-24-09-40-02.bpo-43612.vMGZ4y.rst b/Misc/NEWS.d/next/Library/2021-03-24-09-40-02.bpo-43612.vMGZ4y.rst
@@ -0,0 +1,5 @@
+:func:`zlib.compress` now accepts a wbits parameter which allows users to
+compress data as a raw deflate block without zlib headers and trailers in
+one go. Previously this required instantiating a ``zlib.compressobj``. It
+also provides a faster alternative to ``gzip.compress`` when wbits=31 is
+used.
diff --git a/Modules/clinic/zlibmodule.c.h b/Modules/clinic/zlibmodule.c.h
diff --git a/Modules/zlibmodule.c b/Modules/zlibmodule.c
@@ -205,13 +205,15 @@ zlib.compress
    /
    level: int(c_default="Z_DEFAULT_COMPRESSION") = Z_DEFAULT_COMPRESSION
        Compression level, in 0-9 or -1.
+    wbits: int(c_default="MAX_WBITS") = MAX_WBITS
+        The window buffer size and container format.

 Returns a bytes object containing compressed data.
 [clinic start generated code]*/

 static PyObject *
-zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
-/*[clinic end generated code: output=d80906d73f6294c8 input=638d54b6315dbed3]*/
+zlib_compress_impl(PyObject *module, Py_buffer *data, int level, int wbits)
+/*[clinic end generated code: output=46bd152fadd66df2 input=c4d06ee5782a7e3f]*/
 {
    PyObject *RetVal = NULL;
    Py_ssize_t obuflen = DEF_BUF_SIZE;
@@ -227,7 +229,8 @@ zlib_compress_impl(PyObject *module, Py_buffer *data, int level)
    zst.zalloc = PyZlib_Malloc;
    zst.zfree = PyZlib_Free;
    zst.next_in = ibuf;
-    int err = deflateInit(&zst, level);
+    int err = deflateInit2(&zst, level, DEFLATED, wbits, DEF_MEM_LEVEL,
+                           Z_DEFAULT_STRATEGY);

    switch (err) {
    case Z_OK: