Merge pull request #5504 from radarhere/tuple_size

Added specific error messages when ink has incorrect number of bands
python-pillow · May 23, 2021 · affa059 · affa059
2 parents cc2db17 + a141268
commit affa059
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/Tests/test_image_access.py b/Tests/test_image_access.py
@@ -355,6 +355,24 @@ def test_putpixel_type_error1(self, mode):
             with pytest.raises(TypeError, match="color must be int or tuple"):
                 im.putpixel((0, 0), v)
 
+    @pytest.mark.parametrize(
+        ("mode", "band_numbers", "match"),
+        (
+            ("L", (0, 2), "color must be int or single-element tuple"),
+            ("LA", (0, 3), "color must be int, or tuple of one or two elements"),
+            (
+                "RGB",
+                (0, 2, 5),
+                "color must be int, or tuple of one, three or four elements",
+            ),
+        ),
+    )
+    def test_putpixel_invalid_number_of_bands(self, mode, band_numbers, match):
+        im = hopper(mode)
+        for band_number in band_numbers:
+            with pytest.raises(TypeError, match=match):
+                im.putpixel((0, 0), (0,) * band_number)
+
     @pytest.mark.parametrize("mode", IMAGE_MODES2)
     def test_putpixel_type_error2(self, mode):
         im = hopper(mode)

diff --git a/src/_imaging.c b/src/_imaging.c
@@ -498,7 +498,7 @@ getink(PyObject *color, Imaging im, char *ink) {
        be cast to either UINT8 or INT32 */
 
     int rIsInt = 0;
-    if (PyTuple_Check(color) && PyTuple_Size(color) == 1) {
+    if (PyTuple_Check(color) && PyTuple_GET_SIZE(color) == 1) {
         color = PyTuple_GetItem(color, 0);
     }
     if (im->type == IMAGING_TYPE_UINT8 || im->type == IMAGING_TYPE_INT32 ||
@@ -527,7 +527,10 @@ getink(PyObject *color, Imaging im, char *ink) {
             if (im->bands == 1) {
                 /* unsigned integer, single layer */
                 if (rIsInt != 1) {
-                    if (!PyArg_ParseTuple(color, "L", &r)) {
+                    if (PyTuple_GET_SIZE(color) != 1) {
+                        PyErr_SetString(PyExc_TypeError, "color must be int or single-element tuple");
+                        return NULL;
+                    } else if (!PyArg_ParseTuple(color, "L", &r)) {
                         return NULL;
                     }
                 }
@@ -542,13 +545,20 @@ getink(PyObject *color, Imaging im, char *ink) {
                     g = (UINT8)(r >> 8);
                     r = (UINT8)r;
                 } else {
+                    int tupleSize = PyTuple_GET_SIZE(color);
                     if (im->bands == 2) {
-                        if (!PyArg_ParseTuple(color, "L|i", &r, &a)) {
+                        if (tupleSize != 1 && tupleSize != 2) {
+                            PyErr_SetString(PyExc_TypeError, "color must be int, or tuple of one or two elements");
+                            return NULL;
+                        } else if (!PyArg_ParseTuple(color, "L|i", &r, &a)) {
                             return NULL;
                         }
                         g = b = r;
                     } else {
-                        if (!PyArg_ParseTuple(color, "Lii|i", &r, &g, &b, &a)) {
+                        if (tupleSize != 3 && tupleSize != 4) {
+                            PyErr_SetString(PyExc_TypeError, "color must be int, or tuple of one, three or four elements");
+                            return NULL;
+                        } else if (!PyArg_ParseTuple(color, "Lii|i", &r, &g, &b, &a)) {
                             return NULL;
                         }
                     }