Rewrite axline with custom Line2D subclass.

anntzer · anntzer · commit d83cfc5d645e · 2019-10-22T23:24:22.000+02:00
This makes it also work on non-linear scales.

Also don't use add_line directly but manually copy most of add_line, to
ensure data limits are correctly set.
diff --git a/doc/api/axes_api.rst b/doc/api/axes_api.rst
@@ -87,6 +87,7 @@ Spans
    Axes.axhspan
    Axes.axvline
    Axes.axvspan
+   Axes.axline
 
 Spectral
 --------
diff --git a/doc/users/next_whats_new/2017-12-08-axline.rst b/doc/users/next_whats_new/2017-12-08-axline.rst
@@ -1,5 +1,5 @@
-New `axline` method
--------------------
+New `~.axes.Axes.axline` method
+-------------------------------
 
 A new `~.axes.Axes.axline` method has been added to draw infinitely long lines
 that pass through two points.
diff --git a/examples/subplots_axes_and_figures/axhspan_demo.py b/examples/subplots_axes_and_figures/axhspan_demo.py
@@ -4,32 +4,33 @@
 ============
 
 Create lines or rectangles that span the axes in either the horizontal or
-vertical direction.
+vertical direction, and lines than span the axes with an arbitrary orientation.
 """
+
 import numpy as np
 import matplotlib.pyplot as plt
 
 t = np.arange(-1, 2, .01)
 s = np.sin(2 * np.pi * t)
 
-plt.plot(t, s)
-# Draw a thick red hline at y=0 that spans the xrange
-plt.axhline(linewidth=8, color='#d62728')
-
-# Draw a default hline at y=1 that spans the xrange
-plt.axhline(y=1)
-
-# Draw a default vline at x=1 that spans the yrange
-plt.axvline(x=1)
-
-# Draw a thick blue vline at x=0 that spans the upper quadrant of the yrange
-plt.axvline(x=0, ymin=0.75, linewidth=8, color='#1f77b4')
-
-# Draw a default hline at y=.5 that spans the middle half of the axes
-plt.axhline(y=.5, xmin=0.25, xmax=0.75)
-
-plt.axhspan(0.25, 0.75, facecolor='0.5', alpha=0.5)
-
-plt.axvspan(1.25, 1.55, facecolor='#2ca02c', alpha=0.5)
+fig, ax = plt.subplots()
+
+ax.plot(t, s)
+# Thick red horizontal line at y=0 that spans the xrange.
+ax.axhline(linewidth=8, color='#d62728')
+# Horizontal line at y=1 that spans the xrange.
+ax.axhline(y=1)
+# Vertical line at x=1 that spans the yrange.
+ax.axvline(x=1)
+# Thick blue vertical line at x=0 that spans the upper quadrant of the yrange.
+ax.axvline(x=0, ymin=0.75, linewidth=8, color='#1f77b4')
+# Default hline at y=.5 that spans the middle half of the axes.
+ax.axhline(y=.5, xmin=0.25, xmax=0.75)
+# Infinite black line going through (0, 0) to (1, 1).
+ax.axline((0, 0), (1, 1), color='k')
+# 50%-gray rectangle spanning the axes' width from y=0.25 to y=0.75.
+ax.axhspan(0.25, 0.75, facecolor='0.5')
+# Green rectangle spanning the axes' height from x=1.25 to x=1.55.
+ax.axvspan(1.25, 1.55, facecolor='#2ca02c')
 
 plt.show()
diff --git a/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -816,7 +816,7 @@ def axhline(self, y=0, xmin=0, xmax=1, **kwargs):
         --------
         hlines : Add horizontal lines in data coordinates.
         axhspan : Add a horizontal span (rectangle) across the axis.
-        axhline : Add a line with an arbitrary slope.
+        axline : Add a line with an arbitrary slope.
 
         Examples
         --------
@@ -926,6 +926,10 @@ def axline(self, xy1, xy2, **kwargs):
         """
         Add an infinitely long straight line that passes through two points.
 
+        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.
+
         Parameters
         ----------
         xy1, xy2 : (float, float)
@@ -937,61 +941,41 @@ def axline(self, xy1, xy2, **kwargs):
 
         Other Parameters
         ----------------
-        Valid kwargs are :class:`~matplotlib.lines.Line2D` properties,
-        with the exception of 'transform':
+        **kwargs
+            Valid kwargs are :class:`~matplotlib.lines.Line2D` properties,
+            with the exception of 'transform':
 
-        %(_Line2D_docstr)s
+            %(_Line2D_docstr)s
 
         Examples
         --------
         Draw a thick red line passing through (0, 0) and (1, 1)::
 
             >>> axline((0, 0), (1, 1), linewidth=4, color='r')
 
-
         See Also
         --------
         axhline : for horizontal lines
         axvline : for vertical lines
-
-        Notes
-        -----
-        Currently this method does not work properly with non-linear axes.
         """
-        if not self.get_xscale() == self.get_yscale() == 'linear':
-            raise NotImplementedError('axline() is only supported on '
-                                      'linearly scaled axes')
 
         if "transform" in kwargs:
             raise TypeError("'transform' is not allowed as a kwarg; "
-                            "axline generates its own transform.")
-
+                            "axline generates its own transform")
         x1, y1 = xy1
         x2, y2 = xy2
-        # If x values the same, we have a vertical line
-        if np.allclose(x1, x2):
-            if np.allclose(y1, y2):
-                raise ValueError(
-                    'Cannot draw a line through two identical points '
-                    f'(got x1={x1}, x2={x2}, y1={y1}, y2={y2}).')
-            line = self.axvline(x1, **kwargs)
-            return line
-
-        slope = (y2 - y1) / (x2 - x1)
-        intercept = y1 - (slope * x1)
-
-        xtrans = mtransforms.BboxTransformTo(self.viewLim)
-        viewLimT = mtransforms.TransformedBbox(
-            self.viewLim,
-            mtransforms.Affine2D().rotate_deg(90).scale(-1, 1))
-        ytrans = (mtransforms.BboxTransformTo(viewLimT) +
-                  mtransforms.Affine2D().scale(slope).translate(0, intercept))
-        trans = mtransforms.blended_transform_factory(xtrans, ytrans)
-
-        line = mlines.Line2D([0, 1], [0, 1],
-                             transform=trans + self.transData,
-                             **kwargs)
-        self.add_line(line)
+        line = mlines._AxLine([x1, x2], [y1, y2], **kwargs)
+        # Like add_line, but correctly handling data limits.
+        self._set_artist_props(line)
+        if line.get_clip_path() is None:
+            line.set_clip_path(self.patch)
+        if not line.get_label():
+            line.set_label(f"_line{len(self.lines)}")
+        self.lines.append(line)
+        line._remove_method = self.lines.remove
+        self.update_datalim([xy1, xy2])
+
+        self._request_autoscale_view()
         return line
 
     @docstring.dedent_interpd
diff --git a/lib/matplotlib/lines.py b/lib/matplotlib/lines.py
@@ -15,7 +15,8 @@
     _to_unmasked_float_array, ls_mapper, ls_mapper_r, STEP_LOOKUP_MAP)
 from .markers import MarkerStyle
 from .path import Path
-from .transforms import Bbox, TransformedPath
+from .transforms import (
+    Affine2D, Bbox, BboxTransformFrom, BboxTransformTo, TransformedPath)
 
 # Imported here for backward compatibility, even though they don't
 # really belong.
@@ -1448,6 +1449,47 @@ def is_dashed(self):
         return self._linestyle in ('--', '-.', ':')
 
 
+class _AxLine(Line2D):
+    """
+    A helper class that implements `~.Axes.axline`, by recomputing the artist
+    transform at draw time.
+    """
+
+    def get_transform(self):
+        ax = self.axes
+        (x1, y1), (x2, y2) = ax.transScale.transform([*zip(*self.get_data())])
+        dx = x2 - x1
+        dy = y2 - y1
+        if np.allclose(x1, x2):
+            if np.allclose(y1, y2):
+                raise ValueError(
+                    f"Cannot draw a line through two identical points "
+                    f"(x={self.get_xdata()}, y={self.get_ydata()})")
+            # First send y1 to 0 and y2 to 1.
+            return (Affine2D.from_values(1, 0, 0, 1 / dy, 0, -y1 / dy)
+                    + ax.get_xaxis_transform(which="grid"))
+        if np.allclose(y1, y2):
+            # First send x1 to 0 and x2 to 1.
+            return (Affine2D.from_values(1 / dx, 0, 0, 1, -x1 / dx, 0)
+                    + ax.get_yaxis_transform(which="grid"))
+        (vxlo, vylo), (vxhi, vyhi) = ax.transScale.transform(ax.viewLim)
+        # General case: find intersections with view limits in either
+        # direction, and draw between the middle two points.
+        _, start, stop, _ = sorted([
+            (vxlo, y1 + (vxlo - x1) * dy / dx),
+            (vxhi, y1 + (vxhi - x1) * dy / dx),
+            (x1 + (vylo - y1) * dx / dy, vylo),
+            (x1 + (vyhi - y1) * dx / dy, vyhi),
+        ])
+        return (BboxTransformFrom(Bbox([*zip(*self.get_data())]))
+                + BboxTransformTo(Bbox([start, stop]))
+                + ax.transLimits + ax.transAxes)
+
+    def draw(self, renderer):
+        self._transformed_path = None  # Force regen.
+        super().draw(renderer)
+
+
 class VertexSelector:
     """
     Manage the callbacks to maintain a list of selected vertices for
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -2371,6 +2371,12 @@ def axis(*args, emit=True, **kwargs):
     return gca().axis(*args, emit=emit, **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)
+
+
 # Autogenerated by boilerplate.py.  Do not edit as changes will be lost.
 @docstring.copy(Axes.axvline)
 def axvline(x=0, ymin=0, ymax=1, **kwargs):
diff --git a/tools/boilerplate.py b/tools/boilerplate.py
@@ -194,6 +194,7 @@ def boilerplate_gen():
         'axhline',
         'axhspan',
         'axis',
+        'axline',
         'axvline',
         'axvspan',
         'bar',
