Source code for aerosandbox.tools.pretty_plots.annotation

import matplotlib.pyplot as plt
import matplotlib.transforms as transforms


def hline(
    y,
    linestyle="--",
    color="k",
    text: str = None,
    text_xloc=0.5,
    text_ha="center",
    text_va="bottom",
    text_kwargs=None,
    **kwargs
):  # TODO docs
    if text_kwargs is None:
        text_kwargs = {}
    ax = plt.gca()
    plt.axhline(y=y, ls=linestyle, color=color, **kwargs)
    if text is not None:
        trans = transforms.blended_transform_factory(ax.transAxes, ax.transData)
        plt.annotate(
            text=text,
            xy=(text_xloc, y),
            xytext=(0, 0),
            xycoords=trans,
            textcoords="offset points",
            ha=text_ha,
            va=text_va,
            color=color,
            **text_kwargs
        )


def vline(
    x,
    linestyle="--",
    color="k",
    text: str = None,
    text_yloc=0.5,
    text_ha="right",
    text_va="center",
    text_kwargs=None,
    **kwargs
):  # TODO docs
    if text_kwargs is None:
        text_kwargs = {}
    ax = plt.gca()
    plt.axvline(x=x, ls=linestyle, color=color, **kwargs)
    if text is not None:
        trans = transforms.blended_transform_factory(ax.transData, ax.transAxes)
        plt.annotate(
            text=text,
            xy=(x, text_yloc),
            xytext=(0, 0),
            xycoords=trans,
            textcoords="offset points",
            ha=text_ha,
            va=text_va,
            color=color,
            rotation=90,
            **text_kwargs
        )