Backport PR #18030: Fix PolyCollection.set_verts optimization.

jklymak · meeseeksmachine · commit ba9e2da0c2a2 · 2020-07-23T16:23:20.000Z
diff --git a/lib/matplotlib/collections.py b/lib/matplotlib/collections.py
@@ -1107,7 +1107,7 @@ def __init__(self, verts, sizes=None, closed=True, **kwargs):
         verts : list of array-like
             The sequence of polygons [*verts0*, *verts1*, ...] where each
             element *verts_i* defines the vertices of polygon *i* as a 2D
-            array-like of of shape (M, 2).
+            array-like of shape (M, 2).
         sizes : array-like, default: None
             Squared scaling factors for the polygons. The coordinates of each
             polygon *verts_i* are multiplied by the square-root of the
@@ -1134,7 +1134,7 @@ def set_verts(self, verts, closed=True):
         verts : list of array-like
             The sequence of polygons [*verts0*, *verts1*, ...] where each
             element *verts_i* defines the vertices of polygon *i* as a 2D
-            array-like of of shape (M, 2).
+            array-like of shape (M, 2).
         closed : bool, default: True
             Whether the polygon should be closed by adding a CLOSEPOLY
             connection at the end.
@@ -1149,7 +1149,7 @@ def set_verts(self, verts, closed=True):
             return
 
         # Fast path for arrays
-        if isinstance(verts, np.ndarray):
+        if isinstance(verts, np.ndarray) and len(verts.shape) == 3:
             verts_pad = np.concatenate((verts, verts[:, :1]), axis=1)
             # Creating the codes once is much faster than having Path do it
             # separately each time by passing closed=True.
diff --git a/lib/matplotlib/tests/test_collections.py b/lib/matplotlib/tests/test_collections.py
@@ -615,6 +615,14 @@ def test_collection_set_verts_array():
         assert np.array_equal(ap._vertices, lp._vertices)
         assert np.array_equal(ap._codes, lp._codes)
 
+    verts_tuple = np.empty(10, dtype=object)
+    verts_tuple[:] = [tuple(tuple(y) for y in x) for x in verts]
+    col_arr_tuple = PolyCollection(verts_tuple)
+    assert len(col_arr._paths) == len(col_arr_tuple._paths)
+    for ap, atp in zip(col_arr._paths, col_arr_tuple._paths):
+        assert np.array_equal(ap._vertices, atp._vertices)
+        assert np.array_equal(ap._codes, atp._codes)
+
 
 def test_blended_collection_autolim():
     a = [1, 2, 4]
