multicolored text strings in matplotlib / python

rainbow.py

import matplotlib.pyplot as plt
from matplotlib import transforms

# See https://stackoverflow.com/a/79389836/2683
plt.style.use('classic')

def multicolored_text(x, y, text, cd, **kw):
    """
    Place text at (x, y) with colors for each word specified in the color
    dictionary cd, which maps strings to colors, and must include a 'default'
    key mapping to the default color.

    Based on https://stackoverflow.com/a/9185851/2683, thanks paul ivanov!
    """
    fig = plt.gcf()
    t = plt.gca().transData

    def get_text_width(text, **kw):
        temp = plt.text(0, 0, text, **kw)
        temp.draw(fig.canvas.get_renderer())
        ex = temp.get_window_extent()
        # Need the text width in data coordinates, since that's what x uses.
        width = t.inverted().transform_bbox(ex).width
        temp.remove()
        return width

    ha = kw.pop('ha', 'left')
    match ha:
        case 'left':
            x = x
        case 'center':
            x -= get_text_width(text, **kw) / 2
        case 'right':
            x -= get_text_width(text, **kw)
        case _:
            raise ValueError(f'invalid value for horizontal alignment {ha}')

    for word in text.split(' '):
        c = cd.get(word, cd['default'])
        text = plt.text(x, y, word + ' ', color=c, transform=t, **kw)
        text.draw(fig.canvas.get_renderer())
        ex = text.get_window_extent()
        # Need the text width in points so that this will work with pdf output.
        width = ex.width / fig.dpi * 72  # Can do via fig.dpi_scale_trans?

        t = transforms.offset_copy(text._transform, x=width, units='points', fig=fig)


plt.figure()
text = "One does not simply use multicolored text"
colors = {
    'default': 'black',
    'simply': 'red',
    'multicolored': 'blue',
}
multicolored_text(.5, .05, text, colors, ha='center', size=20)
plt.savefig("multicolored_text.png")

See this stackoverflow answer for context.

Output: