ENH: clip_path keyword for contour and contourf

rcomer · rcomer · commit 458bde016ac2 · 2023-07-13T07:02:11.000+01:00
diff --git a/doc/users/next_whats_new/contour_clip_path.rst b/doc/users/next_whats_new/contour_clip_path.rst
@@ -0,0 +1,24 @@
+Clipping for contour plots
+--------------------------
+
+`~.Axes.contour` and `~.Axes.contourf` now accept the *clip_path* parameter.
+
+.. plot::
+    :include-source: true
+
+    import numpy as np
+    import matplotlib.pyplot as plt
+    import matplotlib.patches as mpatches
+
+    x = y = np.arange(-3.0, 3.01, 0.025)
+    X, Y = np.meshgrid(x, y)
+    Z1 = np.exp(-X**2 - Y**2)
+    Z2 = np.exp(-(X - 1)**2 - (Y - 1)**2)
+    Z = (Z1 - Z2) * 2
+
+    fig, ax = plt.subplots()
+    patch = mpatches.RegularPolygon((0, 0), 5, radius=2,
+                                    transform=ax.transData)
+    ax.contourf(X, Y, Z, clip_path=patch)
+
+    plt.show()
diff --git a/lib/matplotlib/contour.py b/lib/matplotlib/contour.py
@@ -751,7 +751,7 @@ def __init__(self, ax, *args,
                  hatches=(None,), alpha=None, origin=None, extent=None,
                  cmap=None, colors=None, norm=None, vmin=None, vmax=None,
                  extend='neither', antialiased=None, nchunk=0, locator=None,
-                 transform=None, negative_linestyles=None,
+                 transform=None, negative_linestyles=None, clip_path=None,
                  **kwargs):
         """
         Draw contour lines or filled regions, depending on
@@ -805,6 +805,7 @@ def __init__(self, ax, *args,
         super().__init__(
             antialiaseds=antialiased,
             alpha=alpha,
+            clip_path=clip_path,
             transform=transform,
         )
         self.axes = ax
@@ -1870,6 +1871,11 @@ def _initialize_x_y(self, z):
 
     The default is taken from :rc:`contour.algorithm`.
 
+clip_path : `~matplotlib.patches.Patch` or `.Path` or `.TransformedPath`
+    Set the clip path.  See `~matplotlib.artist.Artist.set_clip_path`.
+
+    .. versionadded:: 3.8
+
 data : indexable object, optional
     DATA_PARAMETER_PLACEHOLDER
 
diff --git a/lib/matplotlib/contour.pyi b/lib/matplotlib/contour.pyi
@@ -4,8 +4,10 @@ from matplotlib.axes import Axes
 from matplotlib.collections import Collection, PathCollection
 from matplotlib.colors import Colormap, Normalize
 from matplotlib.font_manager import FontProperties
+from matplotlib.path import Path
+from matplotlib.patches import Patch
 from matplotlib.text import Text
-from matplotlib.transforms import Transform
+from matplotlib.transforms import Transform, TransformedPatchPath, TransformedPath
 from matplotlib.ticker import Locator, Formatter
 
 from numpy.typing import ArrayLike
@@ -99,6 +101,7 @@ class ContourSet(ContourLabeler, Collection):
     negative_linestyles: None | Literal[
         "solid", "dashed", "dashdot", "dotted"
     ] | Iterable[Literal["solid", "dashed", "dashdot", "dotted"]]
+    clip_path: Patch | Path | TransformedPath | TransformedPatchPath | None
     labelTexts: list[Text]
     labelCValues: list[ColorType]
     allkinds: list[np.ndarray]
@@ -145,6 +148,7 @@ class ContourSet(ContourLabeler, Collection):
         negative_linestyles: Literal["solid", "dashed", "dashdot", "dotted"]
         | Iterable[Literal["solid", "dashed", "dashdot", "dotted"]]
         | None = ...,
+        clip_path: Patch | Path | TransformedPath | TransformedPatchPath | None = ...,
         **kwargs
     ) -> None: ...
     def legend_elements(
diff --git a/lib/matplotlib/tests/test_contour.py b/lib/matplotlib/tests/test_contour.py
@@ -9,6 +9,7 @@
 import matplotlib as mpl
 from matplotlib import pyplot as plt, rc_context, ticker
 from matplotlib.colors import LogNorm, same_color
+import matplotlib.patches as mpatches
 from matplotlib.testing.decorators import image_comparison
 import pytest
 
@@ -752,6 +753,14 @@ def test_contour_no_args():
         ax.contour(Z=data)
 
 
+def test_contour_clip_path():
+    fig, ax = plt.subplots()
+    data = [[0, 1], [1, 0]]
+    circle = mpatches.Circle([0.5, 0.5], 0.5, transform=ax.transAxes)
+    cs = ax.contour(data, clip_path=circle)
+    assert cs.get_clip_path() is not None
+
+
 def test_bool_autolevel():
     x, y = np.random.rand(2, 9)
     z = (np.arange(9) % 2).reshape((3, 3)).astype(bool)
