igor/script.py

   1 # Copyright (C) 2012 W. Trevor King <wking@tremily.us>
   2 #
   3 # This file is part of igor.
   4 #
   5 # igor is free software: you can redistribute it and/or modify it under the
   6 # terms of the GNU Lesser General Public License as published by the Free
   7 # Software Foundation, either version 3 of the License, or (at your option) any
   8 # later version.
   9 #
  10 # igor is distributed in the hope that it will be useful, but WITHOUT ANY
  11 # WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
  12 # A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
  13 # details.
  14 #
  15 # You should have received a copy of the GNU Lesser General Public License
  16 # along with igor.  If not, see <http://www.gnu.org/licenses/>.
  17
  18 "Common code for scripts distributed with the `igor` package."
  19
  20 from __future__ import absolute_import
  21 import argparse as _argparse
  22 import logging as _logging
  23 import sys as _sys
  24
  25 try:
  26     import matplotlib as _matplotlib
  27     import matplotlib.pyplot as _matplotlib_pyplot
  28 except ImportError as _matplotlib_import_error:
  29     _matplotlib = None
  30
  31 from . import __version__
  32 from . import LOG as _LOG
  33
  34
  35 class Script (object):
  36     log_levels = [_logging.ERROR, _logging.WARNING, _logging.INFO, _logging.DEBUG]
  37
  38     def __init__(self, description=None, filetype='IGOR Binary Wave (.ibw) file'):
  39         self.parser = _argparse.ArgumentParser(
  40             description=description, version=__version__)
  41         self.parser.add_argument(
  42             '-f', '--infile', metavar='FILE', default='-',
  43             help='input {}'.format(filetype))
  44         self.parser.add_argument(
  45             '-o', '--outfile', metavar='FILE', default='-',
  46             help='file for ASCII output')
  47         self.parser.add_argument(
  48             '-p', '--plot', action='store_const', const=True,
  49             help='use Matplotlib to plot any IGOR waves')
  50         self.parser.add_argument(
  51             '-V', '--verbose', action='count', default=0,
  52             help='increment verbosity')
  53         self._num_plots = 0
  54
  55     def run(self, *args, **kwargs):
  56         args = self.parser.parse_args(*args, **kwargs)
  57         if args.infile == '-':
  58             args.infile = _sys.stdin
  59         if args.outfile == '-':
  60             args.outfile = _sys.stdout
  61         if args.verbose > 1:
  62             log_level = self.log_levels[min(args.verbose-1, len(self.log_levels)-1)]
  63             _LOG.setLevel(log_level)
  64         self._run(args)
  65         self.display_plots()
  66
  67     def _run(self, args):
  68         raise NotImplementedError()
  69
  70     def plot_wave(self, args, wave, title=None):
  71         if not args.plot:
  72             return  # no-op
  73         if not _matplotlib:
  74             raise _matplotlib_import_error
  75         if title is None:
  76             title = wave['wave']['wave_header']['bname']
  77         figure = _matplotlib_pyplot.figure()
  78         axes = figure.add_subplot(1, 1, 1)
  79         axes.set_title(title)
  80         axes.plot(wave['wave']['wData'], 'r.')
  81         self._num_plots += 1
  82
  83     def display_plots(self):
  84         if self._num_plots:
  85             _matplotlib_pyplot.show()