Raise an error if sequence is not flattened

python-pillow · Dec 27, 2021 · e0d5417 · e0d5417
1 parent e9294d8
commit e0d5417
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 9 deletions.
diff --git a/Tests/test_image_putdata.py b/Tests/test_image_putdata.py
@@ -1,6 +1,8 @@
 import sys
 from array import array
 
+import pytest
+
 from PIL import Image
 
 from .helper import assert_image_equal, hopper
@@ -93,3 +95,18 @@ def test_array_F():
     im.putdata(arr)
 
     assert len(im.getdata()) == len(arr)
+
+
+def test_not_flattened():
+    im = Image.new("L", (1, 1))
+    with pytest.raises(TypeError):
+        im.putdata([[0]])
+    with pytest.raises(TypeError):
+        im.putdata([[0]], 2)
+
+    with pytest.raises(TypeError):
+        im = Image.new("I", (1, 1))
+        im.putdata([[0]])
+    with pytest.raises(TypeError):
+        im = Image.new("F", (1, 1))
+        im.putdata([[0]])
diff --git a/src/_imaging.c b/src/_imaging.c
@@ -1494,6 +1494,14 @@ _putdata(ImagingObject *self, PyObject *args) {
         return NULL;
     }
 
+#define set_value_to_item(seq, i) \
+op = PySequence_Fast_GET_ITEM(seq, i); \
+if (PySequence_Check(op)) { \
+    PyErr_SetString(PyExc_TypeError, "sequence must be flattened"); \
+    return NULL; \
+} else { \
+    value = PyFloat_AsDouble(op); \
+}
     if (image->image8) {
         if (PyBytes_Check(data)) {
             unsigned char *p;
@@ -1522,11 +1530,12 @@ _putdata(ImagingObject *self, PyObject *args) {
                 PyErr_SetString(PyExc_TypeError, must_be_sequence);
                 return NULL;
             }
+            double value;
             if (scale == 1.0 && offset == 0.0) {
                 /* Clipped data */
                 for (i = x = y = 0; i < n; i++) {
-                    op = PySequence_Fast_GET_ITEM(seq, i);
-                    image->image8[y][x] = (UINT8)CLIP8((int)PyFloat_AsDouble(op));
+                    set_value_to_item(seq, i);
+                    image->image8[y][x] = (UINT8)CLIP8(value);
                     if (++x >= (int)image->xsize) {
                         x = 0, y++;
                     }
@@ -1535,9 +1544,8 @@ _putdata(ImagingObject *self, PyObject *args) {
             } else {
                 /* Scaled and clipped data */
                 for (i = x = y = 0; i < n; i++) {
-                    PyObject *op = PySequence_Fast_GET_ITEM(seq, i);
-                    image->image8[y][x] =
-                        CLIP8((int)(PyFloat_AsDouble(op) * scale + offset));
+                    set_value_to_item(seq, i);
+                    image->image8[y][x] = CLIP8(value * scale + offset);
                     if (++x >= (int)image->xsize) {
                         x = 0, y++;
                     }
@@ -1555,9 +1563,10 @@ _putdata(ImagingObject *self, PyObject *args) {
         switch (image->type) {
             case IMAGING_TYPE_INT32:
                 for (i = x = y = 0; i < n; i++) {
-                    op = PySequence_Fast_GET_ITEM(seq, i);
+                    double value;
+                    set_value_to_item(seq, i);
                     IMAGING_PIXEL_INT32(image, x, y) =
-                        (INT32)(PyFloat_AsDouble(op) * scale + offset);
+                        (INT32)(value * scale + offset);
                     if (++x >= (int)image->xsize) {
                         x = 0, y++;
                     }
@@ -1566,9 +1575,10 @@ _putdata(ImagingObject *self, PyObject *args) {
                 break;
             case IMAGING_TYPE_FLOAT32:
                 for (i = x = y = 0; i < n; i++) {
-                    op = PySequence_Fast_GET_ITEM(seq, i);
+                    double value;
+                    set_value_to_item(seq, i);
                     IMAGING_PIXEL_FLOAT32(image, x, y) =
-                        (FLOAT32)(PyFloat_AsDouble(op) * scale + offset);
+                        (FLOAT32)(value * scale + offset);
                     if (++x >= (int)image->xsize) {
                         x = 0, y++;
                     }