ENH : add step_where kwarg to fill_between{x}

tacaswell · tacaswell · commit 493c4d1a18f0 · 2015-07-16T00:04:08.000-04:00
Add ability to fill between 'step' plots. Closes #643 and #1709
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -13,7 +13,7 @@
 import matplotlib
 
 import matplotlib.cbook as cbook
-from matplotlib.cbook import _string_to_bool, mplDeprecation
+from matplotlib.cbook import mplDeprecation, STEP_LOOKUP_MAP
 import matplotlib.collections as mcoll
 import matplotlib.colors as mcolors
 import matplotlib.contour as mcontour
@@ -4428,13 +4428,11 @@ def fill(self, *args, **kwargs):
 
     @docstring.dedent_interpd
     def fill_between(self, x, y1, y2=0, where=None, interpolate=False,
+                     step_where=None,
                      **kwargs):
         """
         Make filled polygons between two curves.
 
-        Call signature::
-
-          fill_between(x, y1, y2=0, where=None, **kwargs)
 
         Create a :class:`~matplotlib.collections.PolyCollection`
         filling the regions between *y1* and *y2* where
@@ -4462,9 +4460,12 @@ def fill_between(self, x, y1, y2=0, where=None, interpolate=False,
             end points of the filled region will only occur on explicit
             values in the *x* array.
 
+        step_where : {'pre', 'post', 'mid'}, optional
+            If not None, fill with step logic.
+
 
-        Note
-        ----
+        Notes
+        -----
 
         Additional Keyword args passed on to the
         :class:`~matplotlib.collections.PolyCollection`.
@@ -4516,6 +4517,9 @@ def fill_between(self, x, y1, y2=0, where=None, interpolate=False,
             xslice = x[ind0:ind1]
             y1slice = y1[ind0:ind1]
             y2slice = y2[ind0:ind1]
+            if step_where is not None:
+                step_func = STEP_LOOKUP_MAP[step_where]
+                xslice, y1slice, y2slice = step_func(xslice, y1slice, y2slice)
 
             if not len(xslice):
                 continue
@@ -4576,7 +4580,8 @@ def get_interp_point(ind):
         return collection
 
     @docstring.dedent_interpd
-    def fill_betweenx(self, y, x1, x2=0, where=None, **kwargs):
+    def fill_betweenx(self, y, x1, x2=0, where=None,
+                      step_where=None, **kwargs):
         """
         Make filled polygons between two horizontal curves.
 
@@ -4588,19 +4593,27 @@ def fill_betweenx(self, y, x1, x2=0, where=None, **kwargs):
         filling the regions between *x1* and *x2* where
         ``where==True``
 
-          *y* :
+        Parameters
+        ----------
+        y : array
             An N-length array of the y data
 
-          *x1* :
+        x1 : array
             An N-length array (or scalar) of the x data
 
-          *x2* :
+        x2 : array, optional
             An N-length array (or scalar) of the x data
 
-          *where* :
-             If *None*, default to fill between everywhere.  If not *None*,
-             it is a N length numpy boolean array and the fill will
-             only happen over the regions where ``where==True``
+        where : array, optional
+            If *None*, default to fill between everywhere.  If not *None*,
+            it is a N length numpy boolean array and the fill will
+            only happen over the regions where ``where==True``
+
+        step_where : {'pre', 'post', 'mid'}, optional
+            If not None, fill with step logic.
+
+        Notes
+        -----
 
           *kwargs* :
             keyword args passed on to the
@@ -4610,9 +4623,13 @@ def fill_betweenx(self, y, x1, x2=0, where=None, **kwargs):
 
         %(PolyCollection)s
 
+        Examples
+        --------
+
         .. plot:: mpl_examples/pylab_examples/fill_betweenx_demo.py
 
-        .. seealso::
+        See Also
+        --------
 
             :meth:`fill_between`
                 for filling between two sets of y-values
@@ -4649,6 +4666,9 @@ def fill_betweenx(self, y, x1, x2=0, where=None, **kwargs):
             yslice = y[ind0:ind1]
             x1slice = x1[ind0:ind1]
             x2slice = x2[ind0:ind1]
+            if step_where is not None:
+                step_func = STEP_LOOKUP_MAP[step_where]
+                yslice, x1slice, x2slice = step_func(yslice, x1slice, x2slice)
 
             if not len(yslice):
                 continue
