Created handler for PatchCollection

Faruk Fakih · rcomer · commit 7ed98759583d · 2023-11-22T21:09:13.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
@@ -130,7 +130,7 @@ class HandlerStepPatch(HandlerBase):
         trans: Transform,
     ) -> Sequence[Artist]: ...
 
-class HandlerLineCollection(HandlerLine2D):
+class HandlerLineCollection(HandlerPatch):
     def get_numpoints(self, legend: Legend) -> int: ...
     def create_artists(
         self,
@@ -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,28 @@ 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)
+    _, labels = ax.get_legend_handles_labels()
+    assert len(labels) == 1
+    assert labels[0] == 'my_collection'
+
+    leg = ax.legend()
+    handles = leg.legend_handles
+    assert mpl.colors.same_color(handles[0].get_facecolor(), 'red')
+    assert mpl.colors.same_color(handles[0].get_edgecolor(), 'green')
+    assert handles[0].get_linewidth() == 3
+    np.testing.assert_allclose(handles[0].get_linestyle()[1],
+                               pc.get_linestyle()[0][1])
+
+
 def test_handler_numpoints():
     """Test legend handler with numpoints <= 1."""
     # related to #6921 and PR #8478
