acb1a22b0fec95ac6bc67d84f4f4d886dec8771e
[igor.git] / 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         try:
81             axes.plot(wave['wave']['wData'], 'r.')
82         except ValueError as error:
83             _LOG.error('error plotting {}: {}'.format(title, error))
84             pass
85         self._num_plots += 1
86
87     def display_plots(self):
88         if self._num_plots:
89             _matplotlib_pyplot.show()