Oops, fix up figure/axes/plot argument passing in pypiezo.wiggle.

diff --git a/pypiezo/wiggle.py b/pypiezo/wiggle.py
index c1d16fee6101c582c2a2af608def6a94193261a8..8178d0e8831e97891a65ef197731923c213d6a39 100644 (file)
--- a/pypiezo/wiggle.py
+++ b/pypiezo/wiggle.py
@@ -140,15 +140,15 @@ def _setup_plot(piezo, config, output):
     figure.show()
     _matplotlib_pyplot.draw()
     _matplotlib_pyplot.show()
-    return plot
+    return (figure, axes, plot)
 
-def _update_plot(plot, cycle, data):
+def _update_plot(figure, axes, plot, cycle, data):
     plot[0].set_ydata(data[:,0])
     axes.set_ylim([data.min(), data.max()])
     _matplotlib_pyplot.draw()
 
-def _run_wiggles(piezo, config, plot, output, filename=None, group='/',
-                 keypress_test_mode=False):
+def _run_wiggles(piezo, config, figure, axes, plot, output, filename=None,
+                 group='/', keypress_test_mode=False):
     scan_frequency = config['frequency'] * config['samples']
     cycle = 0
     c = _CheckForKeypress(test_mode=keypress_test_mode)
@@ -163,7 +163,8 @@ def _run_wiggles(piezo, config, plot, output, filename=None, group='/',
                 filename=filename, group=group, config=config,
                 cycle=cycle, data=data)
         if plot:
-            _update_plot(plot=plot, cycle=cycle, data=data)
+            _update_plot(
+                figure=figure, axes=axes, plot=plot, cycle=cycle, data=data)
         cycle += 1
 
 def wiggle_for_interference(
@@ -195,10 +196,14 @@ def wiggle_for_interference(
             output=output)
     if plot:
         interactive = _matplotlib.is_interactive()
-        plot_ = _setup_plot(piezo=piezo, config=config, output=output)
+        figure,axes,plot_ = _setup_plot(
+            piezo=piezo, config=config, output=output)
+    else:
+        figure = axes = plot_ = None
     _run_wiggles(
-        piezo=piezo, config=config, plot=plot_, output=output,
-        filename=filename, group=group, keypress_test_mode=keypress_test_mode)
+        piezo=piezo, config=config, figure=figure, axes=axes, plot=plot_,
+        output=output, filename=filename, group=group,
+        keypress_test_mode=keypress_test_mode)
     if plot:
         _matplotlib.interactive(interactive)
     piezo.last_output[config['axis']] = out[-1,0]