Rewrite axline with custom Line2D subclass.

anntzer · anntzer · commit 68d02b418634 · 2019-09-23T16:40:28.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/lib/matplotlib/axes/_axes.py b/lib/matplotlib/axes/_axes.py
@@ -948,50 +948,29 @@ def axline(self, xy1, xy2, **kwargs):
 
             >>> 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,42 @@ def is_dashed(self):
         return self._linestyle in ('--', '-.', ':')
 
 
+class _AxLine(Line2D):
+    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
@@ -2374,6 +2374,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',
