BUG: quiverkey must set the vector figure attribute

efiring · efiring · commit 2a3bf416e0bd · 2014-02-01T14:49:47.000-10:00
diff --git a/lib/matplotlib/quiver.py b/lib/matplotlib/quiver.py
@@ -303,6 +303,7 @@ def _init(self):
             if self.color is not None:
                 self.vector.set_color(self.color)
             self.vector.set_transform(self.Q.get_transform())
+            self.vector.set_figure(self.get_figure())
             self._initialized = True
 
     def _text_x(self, x):
diff --git a/lib/matplotlib/tests/baseline_images/test_quiver/quiver_with_key_test_image.png b/lib/matplotlib/tests/baseline_images/test_quiver/quiver_with_key_test_image.png
diff --git a/lib/matplotlib/tests/test_quiver.py b/lib/matplotlib/tests/test_quiver.py
@@ -5,24 +5,24 @@
 import sys
 from matplotlib import pyplot as plt
 from matplotlib.testing.decorators import cleanup
+from matplotlib.testing.decorators import image_comparison
 
 
-WRITER_OUTPUT = dict(ffmpeg='mp4', ffmpeg_file='mp4',
-                     mencoder='mp4', mencoder_file='mp4',
-                     avconv='mp4', avconv_file='mp4',
-                     imagemagick='gif', imagemagick_file='gif')
+def draw_quiver(ax):
+    X, Y = np.meshgrid(np.arange(0, 2 * np.pi, 1),
+                       np.arange(0, 2 * np.pi, 1))
+    U = np.cos(X)
+    V = np.sin(Y)
+
+    Q = ax.quiver(U, V)
+    return Q
 
 
 @cleanup
 def test_quiver_memory_leak():
     fig, ax = plt.subplots()
 
-    X, Y = np.meshgrid(np.arange(0, 2 * np.pi, .04),
-                        np.arange(0, 2 * np.pi, .04))
-    U = np.cos(X)
-    V = np.sin(Y)
-
-    Q = ax.quiver(U, V)
+    Q = draw_quiver(ax)
     ttX = Q.X
     Q.remove()
 
@@ -35,20 +35,32 @@ def test_quiver_memory_leak():
 def test_quiver_key_memory_leak():
     fig, ax = plt.subplots()
 
-    X, Y = np.meshgrid(np.arange(0, 2 * np.pi, .04),
-                        np.arange(0, 2 * np.pi, .04))
-    U = np.cos(X)
-    V = np.sin(Y)
-
-    Q = ax.quiver(U, V)
+    Q = draw_quiver(ax)
 
     qk = ax.quiverkey(Q, 0.5, 0.92, 2, r'$2 \frac{m}{s}$',
-                   labelpos='W',
-                   fontproperties={'weight': 'bold'})
+                      labelpos='W',
+                      fontproperties={'weight': 'bold'})
     assert sys.getrefcount(qk) == 3
     qk.remove()
     assert sys.getrefcount(qk) == 2
 
+
+@image_comparison(baseline_images=['quiver_with_key_test_image'],
+                  extensions=['png'])
+def test_quiver_with_key():
+    fig, ax = plt.subplots()
+    ax.margins(0.1)
+
+    Q = draw_quiver(ax)
+
+    qk = ax.quiverkey(Q, 0.5, 0.95, 2,
+                      r'$2\, \mathrm{m}\, \mathrm{s}^{-1}$',
+                      coordinates='figure',
+                      labelpos='W',
+                      fontproperties={'weight': 'bold',
+                                      'size': 'large'})
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()
