From b290d098fee408c7811a1afe507740d70f22f6c4 Mon Sep 17 00:00:00 2001
From: "W. Trevor King" <wking@tremily.us>
Date: Sat, 21 Jul 2012 11:30:12 -0400
Subject: [PATCH] Add -p/--plot option so scripts will plot waves.

---
 bin/igorbinarywave.py       | 21 +++++++++--------
 bin/igorpackedexperiment.py | 45 ++++++++++++++++++++++---------------
 igor/script.py              | 28 +++++++++++++++++++++++
 3 files changed, 67 insertions(+), 27 deletions(-)

diff --git a/bin/igorbinarywave.py b/bin/igorbinarywave.py
index 6215628..70c1bc3 100755
--- a/bin/igorbinarywave.py
+++ b/bin/igorbinarywave.py
@@ -27,13 +27,16 @@ from igor.binarywave import load
 from igor.script import Script
 
 
-def run(args):
-    wave = load(args.infile)
-    numpy.savetxt(args.outfile, wave['wave']['wData'], fmt='%g', delimiter='\t')
-    if args.verbose > 0:
-        wave['wave'].pop('wData')
-        pprint.pprint(wave)
-
-s = Script(description=__doc__)
-s._run = run
+class WaveScript (Script):
+    def _run(self, args):
+        wave = load(args.infile)
+        numpy.savetxt(
+            args.outfile, wave['wave']['wData'], fmt='%g', delimiter='\t')
+        self.plot_wave(args, wave)
+        if args.verbose > 0:
+            wave['wave'].pop('wData')
+            pprint.pprint(wave)
+
+
+s = WaveScript(description=__doc__)
 s.run()
diff --git a/bin/igorpackedexperiment.py b/bin/igorpackedexperiment.py
index 6664a96..af8bafc 100755
--- a/bin/igorpackedexperiment.py
+++ b/bin/igorpackedexperiment.py
@@ -23,25 +23,34 @@ import pprint
 
 import numpy
 
-from igor.packed import load
+from igor.packed import load, walk
+from igor.record.wave import WaveRecord
 from igor.script import Script
 
 
-def run(args):
-    records,filesystem = load(args.infile)
-    if hasattr(args.outfile, 'write'):
-        f = args.outfile  # filename is actually a stream object
-    else:
-        f = open(args.outfile, 'w')
-    try:
-        f.write(pprint.pformat(records))
-        f.write('\n')
-    finally:
-        if f != args.outfile:
-            f.close()
-    if args.verbose > 0:
-        pprint.pprint(filesystem)
-
-s = Script(description=__doc__, filetype='IGOR Packed Experiment (.pxp) file')
-s._run = run
+class PackedScript (Script):
+    def _run(self, args):
+        self.args = args
+        records,filesystem = load(args.infile)
+        if hasattr(args.outfile, 'write'):
+            f = args.outfile  # filename is actually a stream object
+        else:
+            f = open(args.outfile, 'w')
+        try:
+            f.write(pprint.pformat(records))
+            f.write('\n')
+        finally:
+            if f != args.outfile:
+                f.close()
+        if args.verbose > 0:
+            pprint.pprint(filesystem)
+        walk(filesystem, self._plot_wave_callback)
+
+    def _plot_wave_callback(self, dirpath, key, value):
+        if isinstance(value, WaveRecord):
+            self.plot_wave(self.args, value.wave, title=dirpath + [key])
+
+
+s = PackedScript(
+    description=__doc__, filetype='IGOR Packed Experiment (.pxp) file')
 s.run()
diff --git a/igor/script.py b/igor/script.py
index cbfd519..e14db5e 100644
--- a/igor/script.py
+++ b/igor/script.py
@@ -22,6 +22,12 @@ import argparse as _argparse
 import logging as _logging
 import sys as _sys
 
+try:
+    import matplotlib as _matplotlib
+    import matplotlib.pyplot as _matplotlib_pyplot
+except ImportError as _matplotlib_import_error:
+    _matplotlib = None
+
 from . import __version__
 from . import LOG as _LOG
 
@@ -38,9 +44,13 @@ class Script (object):
         self.parser.add_argument(
             '-o', '--outfile', metavar='FILE', default='-',
             help='file for ASCII output')
+        self.parser.add_argument(
+            '-p', '--plot', action='store_const', const=True,
+            help='use Matplotlib to plot any IGOR waves')
         self.parser.add_argument(
             '-V', '--verbose', action='count', default=0,
             help='increment verbosity')
+        self._num_plots = 0
 
     def run(self, *args, **kwargs):
         args = self.parser.parse_args(*args, **kwargs)
@@ -52,6 +62,24 @@ class Script (object):
             log_level = self.log_levels[min(args.verbose-1, len(self.log_levels)-1)]
             _LOG.setLevel(log_level)
         self._run(args)
+        self.display_plots()
 
     def _run(self, args):
         raise NotImplementedError()
+
+    def plot_wave(self, args, wave, title=None):
+        if not args.plot:
+            return  # no-op
+        if not _matplotlib:
+            raise _matplotlib_import_error
+        if title is None:
+            title = wave['wave']['wave_header']['bname']
+        figure = _matplotlib_pyplot.figure()
+        axes = figure.add_subplot(1, 1, 1)
+        axes.set_title(title)
+        axes.plot(wave['wave']['wData'], 'r.')
+        self._num_plots += 1
+
+    def display_plots(self):
+        if self._num_plots:
+            _matplotlib_pyplot.show()
-- 
2.26.2