From 71dc77d7a88d52f4865dad6e99f325139acb22c5 Mon Sep 17 00:00:00 2001 From: Aitik Gupta Date: Thu, 29 Apr 2021 07:08:47 +0530 Subject: [PATCH 1/6] Expand set_array for array-like input --- lib/matplotlib/cm.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py index 5326890ff1fd..f654628938a7 100644 --- a/lib/matplotlib/cm.py +++ b/lib/matplotlib/cm.py @@ -361,12 +361,19 @@ def to_rgba(self, x, alpha=None, bytes=False, norm=True): def set_array(self, A): """ - Set the image array from numpy array *A*. + Set the image array from array-like *A*. Parameters ---------- - A : ndarray or None + A : array-like or None """ + A = cbook.safe_masked_invalid(A, copy=True) if A is not None else None + + if (A is not None and A.dtype != np.uint8 and + not np.can_cast(A.dtype, float, "same_kind")): + raise TypeError("Image data of dtype {} cannot be converted to " + "float".format(A.dtype)) + self._A = A def get_array(self): From 9cc32b0a89a5b7b119effb84381e9b191af8f73e Mon Sep 17 00:00:00 2001 From: Aitik Gupta Date: Thu, 29 Apr 2021 07:08:59 +0530 Subject: [PATCH 2/6] Remove redundant copy --- lib/matplotlib/image.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/lib/matplotlib/image.py b/lib/matplotlib/image.py index 0aa34e467fb2..e755b204ab9d 100644 --- a/lib/matplotlib/image.py +++ b/lib/matplotlib/image.py @@ -1346,8 +1346,7 @@ def make_image(self, renderer, magnification=1.0, unsampled=False): def set_data(self, A): """Set the image array.""" - cm.ScalarMappable.set_array(self, - cbook.safe_masked_invalid(A, copy=True)) + cm.ScalarMappable.set_array(self, A) self.stale = True From c4b9f9fcac6231566cce95df64455fb6a148bade Mon Sep 17 00:00:00 2001 From: Aitik Gupta Date: Thu, 29 Apr 2021 07:25:30 +0530 Subject: [PATCH 3/6] Test set_array for copy and wrong dtype --- lib/matplotlib/tests/test_collections.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py index e226ca131052..c56aa93ed9aa 100644 --- a/lib/matplotlib/tests/test_collections.py +++ b/lib/matplotlib/tests/test_collections.py @@ -677,6 +677,22 @@ def test_collection_set_verts_array(): assert np.array_equal(ap._codes, atp._codes) +def test_collection_set_array(): + vals = [*range(10)] + + # Test set_array with list + c = Collection() + c.set_array(vals) + + # Test set_array with wrong dtype + with pytest.raises(TypeError, match="^Image data of dtype"): + c.set_array("wrong_input") + + # Test if array kwarg is copied + vals[5] = 45 + assert np.not_equal(vals, c.get_array()).any() + + def test_blended_collection_autolim(): a = [1, 2, 4] height = .2 From 1139e6af97ed1780b881fb74eb962d76095279ab Mon Sep 17 00:00:00 2001 From: Aitik Gupta Date: Thu, 29 Apr 2021 14:38:59 +0530 Subject: [PATCH 4/6] Simplify logic Co-authored-by: Tim Hoffmann <2836374+timhoffm@users.noreply.github.com> --- lib/matplotlib/cm.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py index f654628938a7..4b1988f268c2 100644 --- a/lib/matplotlib/cm.py +++ b/lib/matplotlib/cm.py @@ -367,10 +367,12 @@ def set_array(self, A): ---------- A : array-like or None """ - A = cbook.safe_masked_invalid(A, copy=True) if A is not None else None - - if (A is not None and A.dtype != np.uint8 and - not np.can_cast(A.dtype, float, "same_kind")): + if A is None: + self._A = None + return + + A = cbook.safe_masked_invalid(A, copy=True) + if not np.can_cast(A.dtype, float, "same_kind"): raise TypeError("Image data of dtype {} cannot be converted to " "float".format(A.dtype)) From 2044b57c743f9bf1ab93f0693e030bfca43e2235 Mon Sep 17 00:00:00 2001 From: Aitik Gupta Date: Thu, 29 Apr 2021 14:42:53 +0530 Subject: [PATCH 5/6] Flake8 fix --- lib/matplotlib/cm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py index 4b1988f268c2..a66d20633548 100644 --- a/lib/matplotlib/cm.py +++ b/lib/matplotlib/cm.py @@ -370,7 +370,7 @@ def set_array(self, A): if A is None: self._A = None return - + A = cbook.safe_masked_invalid(A, copy=True) if not np.can_cast(A.dtype, float, "same_kind"): raise TypeError("Image data of dtype {} cannot be converted to " From 0ec672fedb231609329e67374f6505c68c0d6b82 Mon Sep 17 00:00:00 2001 From: Aitik Gupta Date: Thu, 29 Apr 2021 16:06:38 +0530 Subject: [PATCH 6/6] use f-string for the error --- lib/matplotlib/cm.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/matplotlib/cm.py b/lib/matplotlib/cm.py index a66d20633548..bf1f03eb4af8 100644 --- a/lib/matplotlib/cm.py +++ b/lib/matplotlib/cm.py @@ -373,8 +373,8 @@ def set_array(self, A): A = cbook.safe_masked_invalid(A, copy=True) if not np.can_cast(A.dtype, float, "same_kind"): - raise TypeError("Image data of dtype {} cannot be converted to " - "float".format(A.dtype)) + raise TypeError(f"Image data of dtype {A.dtype} cannot be " + "converted to float") self._A = A