feat: extend StepPatch set/get methods, tests

andrzejnovak · andrzejnovak · commit 6d177614ecdd · 2020-09-09T17:50:58.000+02:00
diff --git a/lib/matplotlib/patches.py b/lib/matplotlib/patches.py
@@ -1047,10 +1047,13 @@ def _update_path(self):
             codes.append(np.array([Path.MOVETO] + [Path.LINETO]*(len(xy)-1)))
             self._path = Path(np.vstack(verts), np.hstack(codes))
 
+    def get_data(self):
+        """Get `.StepPatch` values and edges."""
+        return self._values, self._edges
+
     def set_data(self, values, edges=None):
         """
-        Set the values and optionally edges of the
-        StepPatch artist.
+        Set `.StepPatch` values and optionally edges.
 
         Parameters
         ----------
@@ -1065,6 +1068,42 @@ def set_data(self, values, edges=None):
         self._update_path()
         self.stale = True
 
+    def set_values(self, values):
+        """
+        Set `.StepPatch` values.
+
+        Parameters
+        ----------
+        values : 1D array-like
+        """
+        self.set_data(values, edges=None)
+
+    def set_edges(self, edges):
+        """
+        Set `.StepPatch` edges.
+
+        Parameters
+        ----------
+        edges : 1D array-like
+        """
+        self.set_data(None, edges=edges)
+
+    def get_baseline(self):
+        """Get `.StepPatch` baseline value."""
+        return self.baseline
+
+    def set_baseline(self, baseline):
+        """
+        Set `.StepPatch` baseline value.
+
+        Parameters
+        ----------
+        baseline : float or None
+        """
+        self.baseline = baseline
+        self._update_path()
+        self.stale = True
+
 
 class Polygon(Patch):
     """A general polygon patch."""
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -1821,12 +1821,18 @@ def test_stairs_update(fig_test, fig_ref):
     fig_test, test_ax = plt.subplots()
     h = test_ax.stairs([1, 2, 3])
     h.set_values([3, 2, 1])
+    h.set_edges(np.arange(4)+2)
     h.set_data([1, 2, 1], np.arange(4)/2)
     h.set_data([1, 2, 3])
     h.set_data(None, np.arange(4))
+    assert np.allclose(h.get_data()[0], np.arange(1, 4))
+    assert np.allclose(h.get_data()[1], np.arange(4))
+    h.set_baseline(-2)
+    assert h.get_baseline() == -2
+
     # # Ref
     fig_ref, ref_ax = plt.subplots()
-    h = ref_ax.stairs([1, 2, 3])
+    h = ref_ax.stairs([1, 2, 3], baseline=-2)
 
 
 @pytest.mark.xfail
@@ -1839,6 +1845,18 @@ def test_stairs_invalid_mismatch():
     plt.stairs([1, 2], [0, 1])
 
 
+@pytest.mark.xfail
+def test_stairs_invalid_update():
+    h = plt.stairs([1, 2], [0, 1])
+    h.set_edges(np.arange(5))
+
+
+@pytest.mark.xfail
+def test_stairs_invalid_update2():
+    h = plt.stairs([1, 2], [0, 1])
+    h.set_edges([1, np.nan, 2])
+
+
 @image_comparison(['test_stairs_options.png'], remove_text=True)
 def test_stairs_options():
     x, y = np.array([1, 2, 3, 4, 5]), np.array([1, 2, 3, 4]).astype(float)
