add ability to specify file and dataset attributes by prompting the user when running fetch.py

1 files changed, 60 insertions, 3 deletions

diff --git a/fetch.py b/fetch.py
index eaeca88..91dfe4a 100755
--- a/fetch.py
+++ b/fetch.py
@@ -25,7 +25,7 @@ import setup
 from lecroy import LeCroyScope
 from config import get_settings
 
-def fetch(filename, nevents):
+def fetch(filename, nevents, runattrs=None):
     """
     Fetch and save waveform traces from the oscilloscope.
 
@@ -87,6 +87,11 @@ def fetch(filename, nevents):
             except ValueError:
                 pass
 
+    if runattrs is not None:
+        for name in runattrs:
+            for key, value in runattrs[name].items():
+                f[name].attrs[key] = value
+
     # start a timer
     time0 = time.time()
 
@@ -147,6 +152,7 @@ def fetch(filename, nevents):
 
 if __name__ == '__main__':
     import optparse
+    import run_setup
 
     usage = "usage: %prog <filename/prefix> [-n] [-r]"
     parser = optparse.OptionParser(usage, version="%prog 0.1.0")
@@ -156,6 +162,8 @@ if __name__ == '__main__':
                       help="number of runs", default=1)
     parser.add_option("--time", action="store_true", dest="time",
                       help="append time string to filename", default=False)
+    parser.add_option("-c", dest="run_config",
+                      help="run configuration dictionary name", default=None)
     (options, args) = parser.parse_args()
 
     if len(args) < 1:
@@ -164,9 +172,58 @@ if __name__ == '__main__':
     if options.nevents < 1 or options.nruns < 1:
         sys.exit("Please specify a number >= 1 for number of events/runs")
 
+    if options.run_config is not None:
+        options.run_config = getattr(run_setup, options.run_config)
+
+        if 'file' not in options.run_config or \
+                'dataset' not in options.run_config:
+            raise AttributeError("run configuration must contain 'file' and 'dataset' keys")
+
+        scope = LeCroyScope(setup.scope_ip, timeout=20.0)
+
+        # clear the output queue
+        scope.clear()
+
+        # get active channels
+        channels = scope.getchannels()
+
+        # close the socket connection
+        del scope
+
+        runattrs = {'/' : {}}
+
+        print '/'
+
+        for key, fmt in options.run_config['file'].items():
+            prompt = '/%s? ' % key
+            while True:
+                try:
+                    runattrs['/'][key] = fmt(raw_input(prompt))
+                except ValueError as e:
+                    print e
+                    continue
+                break
+
+        for name in ['channel%i' % i for i in channels]:
+            runattrs[name] = {}
+
+            print '/' + name
+
+            for key, fmt in options.run_config['dataset'].items():
+                prompt = '/%s.%s? ' % (name, key)
+                while True:
+                    try:
+                        runattrs[name][key] = fmt(raw_input(prompt))
+                    except ValueError as e:
+                        print e
+                        continue
+                    break
+    else:
+        runattrs = None
+
     if options.nruns == 1 and not options.time:
         try:
-            fetch(args[0], options.nevents)
+            fetch(args[0], options.nevents, runattrs)
         except KeyboardInterrupt:
             pass
     else:
@@ -181,6 +238,6 @@ if __name__ == '__main__':
             print '-' * 65
 
             try:
-                fetch(filename, options.nevents)
+                fetch(filename, options.nevents, runattrs)
             except KeyboardInterrupt:
                 break