From a1412681ff07ebca0e0b32c0122aec6fa7839641 Mon Sep 17 00:00:00 2001 From: Andrew Murray Date: Wed, 19 May 2021 20:19:57 +1000 Subject: [PATCH] Added specific error messages when ink has incorrect number of bands --- Tests/test_image_access.py | 18 ++++++++++++++++++ src/_imaging.c | 16 +++++++++++++--- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/Tests/test_image_access.py b/Tests/test_image_access.py index 78d04946e..7b3036979 100644 --- a/Tests/test_image_access.py +++ b/Tests/test_image_access.py @@ -355,6 +355,24 @@ class TestImagePutPixelError(AccessTest): 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 index 72413e235..cde440004 100644 --- a/src/_imaging.c +++ b/src/_imaging.c @@ -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; } }