Added ability to pass memoryview object to ImageDecoder

bigcat88 · Feb 10, 2023 · c7aceff · c7aceff
1 parent 4ab4bfe
commit c7aceff
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/src/decode.c b/src/decode.c
@@ -116,12 +116,24 @@ _dealloc(ImagingDecoderObject *decoder) {
 
 static PyObject *
 _decode(ImagingDecoderObject *decoder, PyObject *args) {
+    PyObject *input_object, *memview = NULL;
+    Py_buffer *memview_buf;
     UINT8 *buffer;
     Py_ssize_t bufsize;
     int status;
     ImagingSectionCookie cookie;
 
-    if (!PyArg_ParseTuple(args, "y#", &buffer, &bufsize)) {
+    if (!PyArg_ParseTuple(args, "O", &input_object)) {
+        return NULL;
+    }
+
+    if (PyMemoryView_Check(input_object)) {
+        memview = PyMemoryView_FromObject(input_object);
+        memview_buf = PyMemoryView_GET_BUFFER(memview);
+        buffer = memview_buf->buf;
+        bufsize = memview_buf->len;
+    }
+    else if (!PyArg_ParseTuple(args, "y#", &buffer, &bufsize)) {
         return NULL;
     }
 
@@ -135,6 +147,7 @@ _decode(ImagingDecoderObject *decoder, PyObject *args) {
         ImagingSectionLeave(&cookie);
     }
 
+    Py_XDECREF(memview);
     return Py_BuildValue("ii", status, decoder->state.errcode);
 }