bpo-43612: Add wbits parameter to zlib.compress

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.10
+      *wbits* parameter added.
 
 .. 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``.

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"""

Misc/NEWS.d/next/Library/2021-03-24-09-40-02.bpo-43612.vMGZ4y.rst

@@ -0,0 +1,5 @@
+``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.

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,7 @@ 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: