matplotlib
diff --git a/‎examples/text_labels_and_annotations/rainbow_text.py
Copy file name to clipboardExpand all lines: examples/text_labels_and_annotations/rainbow_text.py
+16-11Lines changed: 16 additions & 11 deletions b/‎examples/text_labels_and_annotations/rainbow_text.py
Copy file name to clipboardExpand all lines: examples/text_labels_and_annotations/rainbow_text.py
+16-11Lines changed: 16 additions & 11 deletions
@@ -21,35 +21,40 @@
 import matplotlib.pyplot as plt
 from matplotlib import transforms
 
-def rainbow_text(x, y, strings, colors, **kw):
+def rainbow_text(x, y, strings, colors, ax=None, **kw):
     """
-    Take a list of ``strings`` and colors ``colors`` and place them next to each
+    Take a list of ``strings`` and ``colors`` and place them next to each
     other, with text strings[i] being shown in colors[i].
 
     This example shows how to do both vertical and horizontal text, and will
     pass all keyword arguments to plt.text, so you can set the font size,
     family, etc.
+
+    The text will get added to the ``ax`` axes, if provided, otherwise the
+    currently active axes will be used.
     """
-    t = plt.gca().transData
-    fig = plt.gcf()
+    if ax is None:
+        ax = plt.gca()
+    t = ax.transData
+    canvas = ax.figure.canvas
 
-    #horizontal version
+    # horizontal version
     for s,c in zip(strings, colors):
-        text = plt.text(x, y, " " + s + " ", color=c, transform=t, **kw)
-        text.draw(fig.canvas.get_renderer())
+        text = ax.text(x, y, " " + s + " ", color=c, transform=t, **kw)
+        text.draw(canvas.get_renderer())
         ex = text.get_window_extent()
         t = transforms.offset_copy(text._transform, x=ex.width, units='dots')
 
-    #vertical version
+    # vertical version
     for s,c in zip(strings, colors):
-        text = plt.text(x, y, " " + s + " ", color=c, transform=t,
+        text = ax.text(x, y, " " + s + " ", color=c, transform=t,
                 rotation=90, va='bottom', ha='center', **kw)
-        text.draw(fig.canvas.get_renderer())
+        text.draw(canvas.get_renderer())
         ex = text.get_window_extent()
         t = transforms.offset_copy(text._transform, y=ex.height, units='dots')
 
 
-rainbow_text(40, 540, "all unicorns poop rainbows ! ! !".split(), 
+rainbow_text(0, 0, "all unicorns poop rainbows ! ! !".split(),
         ['red', 'cyan', 'brown', 'green', 'blue', 'purple', 'black'],
         size=18)