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(description=description)
  40         self.parser.add_argument(
  41             '--version', action='version',
  42             version='%(prog)s {}'.format(__version__))
  43         self.parser.add_argument(
  44             '-f', '--infile', metavar='FILE', default='-',
  45             help='input {}'.format(filetype))
  46         self.parser.add_argument(
  47             '-o', '--outfile', metavar='FILE', default='-',
  48             help='file for ASCII output')
  49         self.parser.add_argument(
  50             '-p', '--plot', action='store_const', const=True,
  51             help='use Matplotlib to plot any IGOR waves')
  52         self.parser.add_argument(
  53             '-V', '--verbose', action='count', default=0,
  54             help='increment verbosity')
  55         self._num_plots = 0
  56
  57     def run(self, *args, **kwargs):
  58         args = self.parser.parse_args(*args, **kwargs)
  59         if args.infile == '-':
  60             args.infile = _sys.stdin
  61         if args.outfile == '-':
  62             args.outfile = _sys.stdout
  63         if args.verbose > 1:
  64             log_level = self.log_levels[min(args.verbose-1, len(self.log_levels)-1)]
  65             _LOG.setLevel(log_level)
  66         self._run(args)
  67         self.display_plots()
  68
  69     def _run(self, args):
  70         raise NotImplementedError()
  71
  72     def plot_wave(self, args, wave, title=None):
  73         if not args.plot:
  74             return  # no-op
  75         if not _matplotlib:
  76             raise _matplotlib_import_error
  77         if title is None:
  78             title = wave['wave']['wave_header']['bname']
  79         figure = _matplotlib_pyplot.figure()
  80         axes = figure.add_subplot(1, 1, 1)
  81         axes.set_title(title)
  82         try:
  83             axes.plot(wave['wave']['wData'], 'r.')
  84         except ValueError as error:
  85             _LOG.error('error plotting {}: {}'.format(title, error))
  86             pass
  87         self._num_plots += 1
  88
  89     def display_plots(self):
  90         if self._num_plots:
  91             _matplotlib_pyplot.show()