Created handler for PatchCollection

Faruk Fakih · rcomer · commit c7cc153ea881 · 2023-11-18T18:54:57.000Z
diff --git a/doc/users/next_whats_new/patch_collection_handler.rst b/doc/users/next_whats_new/patch_collection_handler.rst
@@ -0,0 +1,24 @@
+Legend handler for PatchCollection objects
+---------------------------------------------------------------------
+
+PatchCollection objects are now supported in legends. The feature can be used as follows:
+
+.. plot::
+    :include-source: true
+
+    import matplotlib.pyplot as plt
+    from matplotlib.collections import PatchCollection
+    from matplotlib.patches import Polygon
+
+    fig, axs = plt.subplots()
+    p1, p2 = Polygon([[0,0],[100,100],[200,0]]), Polygon([[400,0],[500,100],[600,0]])
+    p3, p4 = Polygon([[700,0],[800,100],[900,0]]), Polygon([[1000,0],[1100,100],[1200,0]])
+    p = PatchCollection([p1,p2], label="a", facecolors='red', edgecolors='black')
+    p2 = PatchCollection([p3,p4], label="ab", color='green')
+    axs.add_collection(p, autolim=True)
+    axs.add_collection(p2, autolim=True)
+    axs.set_xlim(right=1200)
+    axs.set_ylim(top=100)
+    axs.legend()
+
+    plt.show()
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -38,7 +38,7 @@
                                 StepPatch)
 from matplotlib.collections import (
     Collection, CircleCollection, LineCollection, PathCollection,
-    PolyCollection, RegularPolyCollection)
+    PolyCollection, PatchCollection, RegularPolyCollection)
 from matplotlib.text import Text
 from matplotlib.transforms import Bbox, BboxBase, TransformedBbox
 from matplotlib.transforms import BboxTransformTo, BboxTransformFrom
@@ -792,6 +792,7 @@ def draw(self, renderer):
         Patch: legend_handler.HandlerPatch(),
         StepPatch: legend_handler.HandlerStepPatch(),
         LineCollection: legend_handler.HandlerLineCollection(),
+        PatchCollection: legend_handler.HandlerPatchCollection(),
         RegularPolyCollection: legend_handler.HandlerRegularPolyCollection(),
         CircleCollection: legend_handler.HandlerCircleCollection(),
         BarContainer: legend_handler.HandlerPatch(
diff --git a/lib/matplotlib/legend_handler.py b/lib/matplotlib/legend_handler.py
@@ -427,6 +427,32 @@ def create_artists(self, legend, orig_handle,
         return [legline]
 
 
+class HandlerPatchCollection(HandlerPatch):
+    """
+    Handler for `.PatchCollection` instances.
+    """
+    def _default_update_prop(self, legend_handle, orig_handle):
+        lw = orig_handle.get_linewidths()[0]
+        dashes = orig_handle._us_linestyles[0]
+        facecolor = orig_handle.get_facecolor()[0]
+        edgecolor = orig_handle.get_edgecolor()[0]
+        legend_handle.set_facecolor(facecolor)
+        legend_handle.set_edgecolor(edgecolor)
+        legend_handle.set_linestyle(dashes)
+        legend_handle.set_linewidth(lw)
+
+    def create_artists(self, legend, orig_handle,
+                       xdescent, ydescent, width, height, fontsize, trans):
+
+        p = self._create_patch(legend, orig_handle,
+                               xdescent, ydescent, width, height, fontsize)
+
+        self.update_prop(p, orig_handle, legend)
+        p.set_transform(trans)
+
+        return [p]
+
+
 class HandlerRegularPolyCollection(HandlerNpointsYoffsets):
     r"""Handler for `.RegularPolyCollection`\s."""
 
diff --git a/lib/matplotlib/legend_handler.pyi b/lib/matplotlib/legend_handler.pyi
@@ -144,6 +144,19 @@ class HandlerLineCollection(HandlerLine2D):
         trans: Transform,
     ) -> Sequence[Artist]: ...
 
+class HandlerPatchCollection(HandlerLine2D):
+    def create_artists(
+        self,
+        legend: Legend,
+        orig_handle: Artist,
+        xdescent: float,
+        ydescent: float,
+        width: float,
+        height: float,
+        fontsize: float,
+        trans: Transform,
+    ) -> Sequence[Artist]: ...
+
 _T = TypeVar("_T", bound=Artist)
 
 class HandlerRegularPolyCollection(HandlerNpointsYoffsets):
diff --git a/lib/matplotlib/tests/test_legend.py b/lib/matplotlib/tests/test_legend.py
@@ -625,6 +625,25 @@ def test_linecollection_scaled_dashes():
         assert oh.get_linestyles()[0] == lh._dash_pattern
 
 
+def test_patch_collection_handler():
+    fig, ax = plt.subplots()
+    pc = mcollections.PatchCollection([
+        plt.Circle((0, 0), radius=1, facecolor='red', edgecolor='green',
+                   linewidth=3, linestyle='--'),
+        plt.Rectangle((0.5, 0.5), 1, 1),
+    ], match_original=True, label='my_collection')
+
+    ax.add_collection(pc)
+    handles, labels = ax.get_legend_handles_labels()
+    assert len(labels) == 1
+    assert labels[0] == 'my_collection'
+    assert mpl.colors.same_color(handles[0].get_facecolor()[0], 'red')
+    assert mpl.colors.same_color(handles[0].get_edgecolor()[0], 'green')
+    assert handles[0].get_linewidth()[0] == 3
+    np.testing.assert_allclose(handles[0].get_linestyle()[0][1],
+                               pc.get_linestyle()[0][1])
+
+
 def test_handler_numpoints():
     """Test legend handler with numpoints <= 1."""
     # related to #6921 and PR #8478
