Create axline() using slope

timhoffm · timhoffm · commit c65e83642108 · 2020-01-27T20:20:21.000+01:00
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -904,18 +904,26 @@ def axvline(self, x=0, ymin=0, ymax=1, **kwargs):
         return l
 
     @docstring.dedent_interpd
-    def axline(self, xy1, xy2, **kwargs):
+    def axline(self, xy1, xy2=None, *, slope=None, **kwargs):
         """
-        Add an infinitely long straight line that passes through two points.
+        Add an infinitely long straight line.
+
+        The line can be defined either by two points *xy1* and *xy2*, or
+        by one point *xy1* and a *slope*.
 
         This draws a straight line "on the screen", regardless of the x and y
         scales, and is thus also suitable for drawing exponential decays in
-        semilog plots, power laws in loglog plots, etc.
+        semilog plots, power laws in loglog plots, etc. However, *slope*
+        can only be used with linear scales; for all other scales, please
+        specify the line using the points *xy1*, *xy2*.
 
         Parameters
         ----------
         xy1, xy2 : (float, float)
             Points for the line to pass through.
+            Either *xy2* or *slope* has to be given.
+        slope : float, optional
+            The slope of the line. Either *xy2* or *slope* has to be given.
 
         Returns
         -------
@@ -940,12 +948,29 @@ def axline(self, xy1, xy2, **kwargs):
         axhline : for horizontal lines
         axvline : for vertical lines
         """
+        def _to_points(xy1, xy2, slope):
+            """
+            Check for a valid combination of input parameters and convert
+            to two points, if necessary.
+            """
+            if (xy2 is None and slope is None or
+                    xy2 is not None and slope is not None):
+                raise TypeError(
+                    "Exactly one of 'xy2' and 'slope' must be given")
+            if xy2 is None:
+                x1, y1 = xy1
+                xy2 = (x1, y1 + 1) if np.isinf(slope) else (x1 + 1, y1 + slope)
+            return xy1, xy2
 
         if "transform" in kwargs:
             raise TypeError("'transform' is not allowed as a kwarg; "
                             "axline generates its own transform")
-        x1, y1 = xy1
-        x2, y2 = xy2
+        if slope is not None and (self.get_xscale() != 'linear' or
+                                  self.get_yscale() != 'linear'):
+            raise TypeError("'slope' cannot be used with non-linear scales")
+
+        datalim = [xy1] if xy2 is None else [xy1, xy2]
+        (x1, y1), (x2, y2) = _to_points(xy1, xy2, slope)
         line = mlines._AxLine([x1, x2], [y1, y2], **kwargs)
         # Like add_line, but correctly handling data limits.
         self._set_artist_props(line)
@@ -955,7 +980,7 @@ def axline(self, xy1, xy2, **kwargs):
             line.set_label(f"_line{len(self.lines)}")
         self.lines.append(line)
         line._remove_method = self.lines.remove
-        self.update_datalim([xy1, xy2])
+        self.update_datalim(datalim)
 
         self._request_autoscale_view()
         return line
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2208,8 +2208,8 @@ def axis(*args, emit=True, **kwargs):
 
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
 @docstring.copy(Axes.axline)
-def axline(xy1, xy2, **kwargs):
-    return gca().axline(xy1, xy2, **kwargs)
+def axline(xy1, xy2=None, *, slope=None, **kwargs):
+    return gca().axline(xy1, xy2=xy2, slope=slope, **kwargs)
 
 
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
diff --git a/lib/matplotlib/tests/test_axes.py b/lib/matplotlib/tests/test_axes.py
@@ -3833,12 +3833,36 @@ def test_axline(fig_test, fig_ref):
     ax.axline((0, 0), (1, 1))
     ax.axline((0, 0), (1, 0), color='C1')
     ax.axline((0, 0.5), (1, 0.5), color='C2')
+    # slopes
+    ax.axline((-0.7, -0.5), slope=0, color='C3')
+    ax.axline((1, -0.5), slope=-0.5, color='C4')
+    ax.axline((-0.5, 1), slope=float('inf'), color='C5')
 
     ax = fig_ref.subplots()
     ax.set(xlim=(-1, 1), ylim=(-1, 1))
     ax.plot([-1, 1], [-1, 1])
     ax.axhline(0, color='C1')
     ax.axhline(0.5, color='C2')
+    # slopes
+    ax.axhline(-0.5, color='C3')
+    ax.plot([-1, 1], [0.5, -0.5], color='C4')
+    ax.axvline(-0.5, color='C5')
+
+
+def test_axline_args():
+    """Exactly one of *xy2* and *slope* must be specified."""
+    fig, ax = plt.subplots()
+    with pytest.raises(TypeError):
+        ax.axline((0, 0))  # missing second parameter
+    with pytest.raises(TypeError):
+        ax.axline((0, 0), (1, 1), slope=1)  # redundant parameters
+    ax.set_xscale('log')
+    with pytest.raises(TypeError):
+        ax.axline((0, 0), slope=1)
+    ax.set_xscale('linear')
+    ax.set_yscale('log')
+    with pytest.raises(TypeError):
+        ax.axline((0, 0), slope=1)
 
 
 @image_comparison(['vlines_basic', 'vlines_with_nan', 'vlines_masked'],
