Directional Arrow Callout Annotations in Matplotlib

matplotlib_arrow_text_boxes.py

"""
This script draws arrows pointing in four directions from a text box using matplotlib.
"""

import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrow
from matplotlib.text import Text

def draw_arrow(ax, x, y, txt, direction, length, width):
    """
    Draw an arrow on a matplotlib axes instance.

    Parameters:
        ax (object): matplotlib axes object
        txt (str): text to display
        direction (str): direction of the arrow ('up', 'down', 'left', 'right')
        length (float): length of the arrow
        width (float): width of the arrow
        x (float): x-coordinate of the text box center
        y (float): y-coordinate of the text box center
    """
    color = '#696969'

    txt = Text(x, y, txt,
               bbox=dict(boxstyle='round,pad=0.9,rounding_size=1.47', fc=color, ec=color),
               size="large", color="w",
               ha="center", va="center")

    ax.add_artist(txt)

    points = ax.transData.inverted().transform(txt.get_window_extent().get_points())
    x_tail, y_tail = points[:, 0].mean(), points[:, 1].mean()

    starting_point_adjustment_factor = 0.05  # this helps the arrow-head shape look properly connected to the bubble
    if direction == 'up':
        y_tail = points[:, 1].max() - (points[:, 1].max() - points[:, 1].min()) * starting_point_adjustment_factor
        dx, dy = 0, length
    elif direction == 'down':
        y_tail = points[:, 1].min() + (points[:, 1].max() - points[:, 1].min()) * starting_point_adjustment_factor
        dx, dy = 0, -length
    elif direction == 'left':
        x_tail = points[:, 0].min() + (points[:, 0].max() - points[:, 0].min()) * starting_point_adjustment_factor
        dx, dy = -length, 0
    elif direction == 'right':
        x_tail = points[:, 0].max() - (points[:, 0].max() - points[:, 0].min()) * starting_point_adjustment_factor
        dx, dy = length, 0
    else:
        raise ValueError("Invalid direction. Choose from 'up', 'down', 'left', 'right'.")

    arrow = FancyArrow(x_tail, y_tail, dx, dy,
                       head_width=width, head_length=length, length_includes_head=True,
                       overhang=0, fc=color, ec=color)

    ax.add_patch(arrow)


def main():
    plt.style.use('dark_background')
    fig, ax = plt.subplots()

    toast_messages = [
        {"text": "short", "x": 0.308, "y": 0.728, "length": 0.1, "width": 0.05, "direction": 'up'},
        {"text": "and this is a longer message", "x": 0.66, "y": 0.5, "length": 0.1, "width": 0.1, "direction": 'right'},
        {"text": "multi-line\nmessage\nwow", "x": 0.236, "y": 0.276, "length": 0.06, "width": 0.18, "direction": 'down'},
    ]

    for msg in toast_messages:
        draw_arrow(ax, msg["x"], msg["y"], msg["text"], msg["direction"], msg["length"], msg["width"])

    ax.set(xlim=(0, 1), ylim=(0, 1))
    fig.tight_layout()
    fig.set_facecolor('#0D1117')
    ax.set_facecolor('#161B21')
    plt.show()


if __name__ == "__main__":
    main()