fix(crc32c): accept any contiguous bytes-like object in value() and extend()

packages/google-crc32c/src/google_crc32c/__init__.py

@@ -34,7 +34,14 @@
 
 extend = impl.extend
 value = impl.value
+GIL_RELEASE_THRESHOLD = impl.GIL_RELEASE_THRESHOLD
 
 Checksum = impl.Checksum
 
-__all__ = ["extend", "value", "Checksum", "implementation"]
+__all__ = [
+    "extend",
+    "value",
+    "Checksum",
+    "GIL_RELEASE_THRESHOLD",
+    "implementation",
+]

packages/google-crc32c/src/google_crc32c/_checksum.py

@@ -21,7 +21,7 @@ class CommonChecksum(object):
     This class should not be used directly and requires an update implementation.
 
     Args:
-        initial_value (Optional[bytes]): the initial chunk of data from
+        initial_value (Optional[bytes-like]): the initial chunk of data from
             which the CRC32C checksum is computed.  Defaults to b''.
     """
 
@@ -36,7 +36,7 @@ def update(self, data):
         """Update the checksum with a new chunk of data.
 
         Args:
-            chunk (Optional[bytes]): a chunk of data used to extend
+            chunk (Optional[bytes-like]): a chunk of data used to extend
                 the CRC32C checksum.
         """
         raise NotImplementedError

packages/google-crc32c/src/google_crc32c/_crc32c.c

@@ -6,35 +6,54 @@
 static const Py_ssize_t gil_threshold = 1024 * 1024;
 
 static int
-_should_release_gil(Py_ssize_t length, PyObject *chunk_obj)
+_should_release_gil(const Py_buffer *chunk)
 {
-    /* Checks if the chunk is immutable (bytes) to prevent concurrent modification,
-     * and large enough to benefit from releasing the GIL. */
-    return (length >= gil_threshold && PyBytes_Check(chunk_obj));
+    /* Release the GIL only for chunks of at least gil_threshold bytes backed
+     * by an immutable bytes object (a bytes instance or a memoryview of one). */
+    PyObject *owner = chunk->obj;
+
+    if (chunk->len < gil_threshold || owner == NULL) {
+        return 0;
+    }
+    if (PyMemoryView_Check(owner)) {
+        owner = PyMemoryView_GET_BASE(owner);
+    }
+    return owner != NULL && PyBytes_Check(owner);
 }
 
-static PyObject *
-_crc32c_extend(PyObject *self, PyObject *args)
+static uint32_t
+_extend_buffer(uint32_t crc, const Py_buffer *chunk)
 {
-    unsigned long crc_input;
-    uint32_t crc;
-    const char *chunk;
-    Py_ssize_t length;
     PyThreadState *save = NULL;
 
-    if (!PyArg_ParseTuple(args, "ky#", &crc_input, &chunk, &length))
-        return NULL;
-
-    if (_should_release_gil(length, PyTuple_GET_ITEM(args, 1))) {
+    if (_should_release_gil(chunk)) {
         save = PyEval_SaveThread();
     }
 
-    crc = crc32c_extend((uint32_t)crc_input, (const uint8_t*)chunk, length);
+    crc = crc32c_extend(crc, (const uint8_t*)chunk->buf, (size_t)chunk->len);
 
     if (save) {
         PyEval_RestoreThread(save);
     }
 
+    return crc;
+}
+
+static PyObject *
+_crc32c_extend(PyObject *self, PyObject *args)
+{
+    unsigned long crc_input;
+    uint32_t crc;
+    Py_buffer chunk;
+
+    /* "y*" accepts any C-contiguous object supporting the buffer protocol
+     * (bytes, bytearray, memoryview, array.array, mmap, ...). */
+    if (!PyArg_ParseTuple(args, "ky*", &crc_input, &chunk))
+        return NULL;
+
+    crc = _extend_buffer((uint32_t)crc_input, &chunk);
+    PyBuffer_Release(&chunk);
+
     return PyLong_FromUnsignedLong(crc);
 }
 
@@ -43,22 +62,14 @@ static PyObject *
 _crc32c_value(PyObject *self, PyObject *args)
 {
     uint32_t crc;
-    const char *chunk;
-    Py_ssize_t length;
-    PyThreadState *save = NULL;
+    Py_buffer chunk;
 
-    if (!PyArg_ParseTuple(args, "y#", &chunk, &length))
+    if (!PyArg_ParseTuple(args, "y*", &chunk))
         return NULL;
 
-    if (_should_release_gil(length, PyTuple_GET_ITEM(args, 0))) {
-        save = PyEval_SaveThread();
-    }
-
-    crc = crc32c_value((const uint8_t*)chunk, length);
-
-    if (save) {
-        PyEval_RestoreThread(save);
-    }
+    /* crc32c_value(data, n) is crc32c_extend(0, data, n). */
+    crc = _extend_buffer(0, &chunk);
+    PyBuffer_Release(&chunk);
 
     return PyLong_FromUnsignedLong(crc);
 }
@@ -86,5 +97,13 @@ static struct PyModuleDef crc32cmodule = {
 PyMODINIT_FUNC
 PyInit__crc32c(void)
 {
-    return PyModule_Create(&crc32cmodule);
+    PyObject *m = PyModule_Create(&crc32cmodule);
+    if (m == NULL) {
+        return NULL;
+    }
+    if (PyModule_AddIntConstant(m, "GIL_RELEASE_THRESHOLD", (long)gil_threshold) < 0) {
+        Py_DECREF(m);
+        return NULL;
+    }
+    return m;
 }

packages/google-crc32c/src/google_crc32c/cext.py

@@ -17,6 +17,7 @@
 import google_crc32c.__config__  # noqa: F401
 from google_crc32c._checksum import CommonChecksum
 from google_crc32c._crc32c import (  # type: ignore
+    GIL_RELEASE_THRESHOLD,
     extend,
     value,
 )
@@ -26,8 +27,10 @@ class Checksum(CommonChecksum):
     """Hashlib-alike helper for CRC32C operations.
 
     Args:
-        initial_value (Optional[bytes]): the initial chunk of data from
-            which the CRC32C checksum is computed.  Defaults to b''.
+        initial_value (Optional[bytes-like]): the initial chunk of data from
+            which the CRC32C checksum is computed.  Any C-contiguous object
+            supporting the buffer protocol (``bytes``, ``bytearray``,
+            ``memoryview``, ...) is accepted.  Defaults to b''.
     """
 
     __slots__ = ("_crc",)
@@ -39,7 +42,10 @@ def update(self, chunk):
         """Update the checksum with a new chunk of data.
 
         Args:
-            chunk (Optional[bytes]): a chunk of data used to extend
+            chunk (Optional[bytes-like]): a chunk of data used to extend
                 the CRC32C checksum.
         """
         self._crc = extend(self._crc, chunk)
+
+
+__all__ = ["extend", "value", "Checksum", "GIL_RELEASE_THRESHOLD"]

packages/google-crc32c/src/google_crc32c/python.py

@@ -16,14 +16,17 @@
 
 from google_crc32c._checksum import CommonChecksum
 
+GIL_RELEASE_THRESHOLD = 1024 * 1024
+
 
 def extend(crc, chunk):
     """Update an existing CRC checksum with new chunk of data.
 
     Args
         crc (int): An existing CRC check sum.
-        chunk (Union[bytes, List[int], Tuple[int]]): A new chunk of data.
-            Intended to be a byte string or similar.
+        chunk (Union[bytes-like, Iterable[int]]): A new chunk of data.
+            Bytes-like objects are hashed as raw bytes; an iterable of
+            ints in ``range(256)`` is also accepted.
 
     Returns
         int: New CRC checksum computed by extending existing CRC
@@ -39,8 +42,9 @@ def value(chunk):
     """Compute a CRC checksum for a chunk of data.
 
     Args
-        chunk (Union[bytes, List[int], Tuple[int]]): A new chunk of data.
-            Intended to be a byte string or similar.
+        chunk (Union[bytes-like, Iterable[int]]): A new chunk of data.
+            Bytes-like objects are hashed as raw bytes; an iterable of
+            ints in ``range(256)`` is also accepted.
 
     Returns
         int: New CRC checksum computed for ``chunk``.
@@ -54,31 +58,38 @@ class Checksum(CommonChecksum):
     """Hashlib-alike helper for CRC32C operations.
 
     Args:
-        initial_value (Optional[bytes]): the initial chunk of data from
+        initial_value (Optional[bytes-like]): the initial chunk of data from
             which the CRC32C checksum is computed.  Defaults to b''.
     """
 
     def __init__(self, initial_value=b""):
         self._crc = 0
-        if initial_value != b"":
-            self.update(initial_value)
+        self.update(initial_value)
 
     def update(self, data):
         """Update the checksum with a new chunk of data.
 
         Args:
-            chunk (Optional[bytes]): a chunk of data used to extend
+            data (Optional[bytes-like]): a chunk of data used to extend
                 the CRC32C checksum.
         """
-        if not isinstance(data, array.array) or data.itemsize != 1:
-            buffer = array.array("B", data)
-        else:
-            buffer = data
-        self._crc = self._crc ^ 0xFFFFFFFF
-        for b in buffer:
-            table_poly = _TABLE[(b ^ self._crc) & 0xFF]
-            self._crc = table_poly ^ ((self._crc >> 8) & 0xFFFFFFFF)
-        self._crc = self._crc ^ 0xFFFFFFFF
+        try:
+            view = memoryview(data)
+        except TypeError:
+            # Not a buffer: an iterable of ints in range(256).
+            view = memoryview(array.array("B", data))
+        with view:
+            # Hash the raw bytes, like the C extension.
+            if not view.c_contiguous:
+                raise BufferError("memoryview: underlying buffer is not C-contiguous")
+            if not view.nbytes:
+                return  # Nothing to hash; cast() rejects empty N-D views.
+            with view.cast("B") as buffer:
+                self._crc = self._crc ^ 0xFFFFFFFF
+                for b in buffer:
+                    table_poly = _TABLE[(b ^ self._crc) & 0xFF]
+                    self._crc = table_poly ^ ((self._crc >> 8) & 0xFFFFFFFF)
+                self._crc = self._crc ^ 0xFFFFFFFF
 
 
 # fmt:off