pyplot: draw a line through 1 or 2 reference points (works with scrolling)

axline.py

import matplotlib.pyplot as plt

def axline(xdata, ydata, ax = plt.subplot(111), *args, **kwargs):
    
    try:
        x1 = xdata[0]
        y1 = ydata[0]
    except:
        # assume first point coordinates given as numbers
        x1 = float(xdata)
        y1 = float(ydata)
    
    try:
        x2 = xdata[1]
        y2 = ydata[1]
    except:
        # assume no second point given - line passes (0, 0)
        x2 = 0.
        y2 = 0.
    
    dx = float(x1 - x2)
    dy = float(y1 - y2)
    
    if dx == 0. and dy == 0.:
        pass
    elif dx == 0.:
        ax.axvline(y1, *args, **kwargs)
    elif dy == 0.:
        ax.axhline(x1, *args, **kwargs)
    else:
        slope = dy / dx
        intercept = y1 - slope * x1

        def calculate_points():
            (x1, x2) = ax.get_xlim()
            return [x1, x2], [slope * x1 + intercept, slope * x2 + intercept]

        line, = ax.plot(*(calculate_points() + args), **kwargs)

        def on_change(axes):
            line.set_data(*calculate_points())

        ax.callbacks.connect('xlim_changed', on_change)
        ax.callbacks.connect('ylim_changed', on_change)