1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function
8import os
9import sys
10import re
11import argparse
12import logging
13import textwrap
15from pyrocko import util, progress
16from pyrocko.squirrel import error
19logger = logging.getLogger('psq.tool.common')
21help_time_format = 'Format: ```YYYY-MM-DD HH:MM:SS.FFF```, truncation allowed.'
24def unwrap(s):
25 if s is None:
26 return None
27 s = s.strip()
28 parts = re.split(r'\n{2,}', s)
29 lines = []
30 for part in parts:
31 plines = part.splitlines()
32 if not any(line.startswith(' ') for line in plines):
33 lines.append(' '.join(plines))
34 else:
35 lines.extend(plines)
37 lines.append('')
39 return '\n'.join(lines)
42def wrap(s):
43 lines = []
44 parts = re.split(r'\n{2,}', s)
45 for part in parts:
46 plines = part.splitlines()
47 if part.startswith('usage:') \
48 or all(line.startswith(' ') for line in plines):
49 lines.extend(plines)
50 else:
51 for line in plines:
52 if not line:
53 lines.append(line)
54 if not line.startswith(' '):
55 lines.extend(
56 textwrap.wrap(line, 79,))
57 else:
58 lines.extend(
59 textwrap.wrap(line, 79, subsequent_indent=' '*24))
61 lines.append('')
63 return '\n'.join(lines)
66def wrap_usage(s):
67 lines = []
68 for line in s.splitlines():
69 if not line.startswith('usage:'):
70 lines.append(line)
71 else:
72 lines.extend(textwrap.wrap(line, 79, subsequent_indent=' '*24))
74 return '\n'.join(lines)
77def formatter_with_width(n):
78 class PyrockoHelpFormatter(argparse.RawDescriptionHelpFormatter):
79 def __init__(self, *args, **kwargs):
80 kwargs['width'] = n
81 kwargs['max_help_position'] = 24
82 argparse.RawDescriptionHelpFormatter.__init__(
83 self, *args, **kwargs)
85 # fix alignment problems, with the post-processing wrapping
86 self._action_max_length = 24
88 return PyrockoHelpFormatter
91class PyrockoArgumentParser(argparse.ArgumentParser):
93 # We want to convert the --help outputs to rst for the html docs. Problem
94 # is that argparse's HelpFormatters to date have no public interface which
95 # we could use to achieve this. The solution here is a bit clunky but works
96 # ok for Squirrel. We allow markup like ``code`` which is kept when
97 # producing rst (by parsing the final --help output) but stripped out when
98 # doing normal --help. This leads to a problem with the internal wrapping
99 # of argparse does this before the stripping. To solve, we render with
100 # argparse to a very wide width and do the wrapping in postprocessing.
101 # ``code`` is replaced with just code in normal output. ```code``` is
102 # replaced with 'code' in normal output and with ``code`` rst output. rst
103 # output is selected with environment variable PYROCKO_RST_HELP=1.
104 # The script maintenance/argparse_help_to_rst.py extracts the rst help
105 # and generates the rst files for the docs.
107 def __init__(
108 self, prog=None, usage=None, description=None, epilog=None,
109 **kwargs):
111 kwargs['formatter_class'] = formatter_with_width(1000000)
113 description = unwrap(description)
114 epilog = unwrap(epilog)
116 argparse.ArgumentParser.__init__(
117 self, prog=prog, usage=usage, description=description,
118 epilog=epilog, **kwargs)
120 if hasattr(self, '_action_groups'):
121 for group in self._action_groups:
122 if group.title == 'positional arguments':
123 group.title = 'Positional arguments'
125 elif group.title == 'optional arguments':
126 group.title = 'Optional arguments'
128 elif group.title == 'options':
129 group.title = 'Options'
131 self.raw_help = False
133 def format_help(self, *args, **kwargs):
134 s = argparse.ArgumentParser.format_help(self, *args, **kwargs)
136 # replace usage with wrapped one from argparse because naive wrapping
137 # does not look good.
138 formatter_class = self.formatter_class
139 self.formatter_class = formatter_with_width(79)
140 usage = self.format_usage()
141 self.formatter_class = formatter_class
143 lines = []
144 for line in s.splitlines():
145 if line.startswith('usage:'):
146 lines.append(usage)
147 else:
148 lines.append(line)
150 s = '\n'.join(lines)
152 if os.environ.get('PYROCKO_RST_HELP', '0') == '0':
153 s = s.replace('```', '\'')
154 s = s.replace('``', '')
155 s = wrap(s)
156 else:
157 s = s.replace('```', '``')
158 s = wrap_usage(s)
160 return s
163class SquirrelArgumentParser(PyrockoArgumentParser):
164 '''
165 Parser for CLI arguments with a some extras for Squirrel based apps.
167 :param command:
168 Implementation of the command.
169 :type command:
170 :py:class:`SquirrelCommand` (or module providing the same interface).
172 :param subcommands:
173 Implementations of subcommands.
174 :type subcommands:
175 :py:class:`list` of :py:class:`SquirrelCommand` (or modules providing
176 the same interface).
178 :param \\*args:
179 Handed through to base class's init.
181 :param \\*\\*kwargs:
182 Handed through to base class's init.
183 '''
185 def __init__(self, *args, command=None, subcommands=[], **kwargs):
187 self._command = command
188 self._subcommands = subcommands
189 self._have_selection_arguments = False
190 self._have_query_arguments = False
192 kwargs['add_help'] = False
193 PyrockoArgumentParser.__init__(self, *args, **kwargs)
194 add_standard_arguments(self)
195 self._command = None
196 self._subcommands = []
197 if command:
198 self.set_command(command)
200 if subcommands:
201 self.set_subcommands(subcommands)
203 def set_command(self, command):
204 command.setup(self)
205 self.set_defaults(target=command.run)
207 def set_subcommands(self, subcommands):
208 subparsers = self.add_subparsers(
209 metavar='SUBCOMMAND',
210 title='Subcommands')
212 for mod in subcommands:
213 subparser = mod.make_subparser(subparsers)
214 if subparser is None:
215 raise Exception(
216 'make_subparser(subparsers) must return the created '
217 'parser.')
219 mod.setup(subparser)
220 subparser.set_defaults(target=mod.run, subparser=subparser)
222 def parse_args(self, args=None, namespace=None):
223 '''
224 Parse arguments given on command line.
226 Extends the functionality of
227 :py:meth:`argparse.ArgumentParser.parse_args` to process and handle the
228 standard options ``--loglevel``, ``--progress`` and ``--help``.
229 '''
231 args = PyrockoArgumentParser.parse_args(
232 self, args=args, namespace=namespace)
234 eff_parser = args.__dict__.get('subparser', self)
236 process_standard_arguments(eff_parser, args)
238 if eff_parser._have_selection_arguments:
239 def make_squirrel():
240 return squirrel_from_selection_arguments(args)
242 args.make_squirrel = make_squirrel
244 if eff_parser._have_query_arguments:
245 try:
246 args.squirrel_query = squirrel_query_from_arguments(args)
247 except (error.SquirrelError, error.ToolError) as e:
248 logger.fatal(str(e))
249 sys.exit(1)
251 return args
253 def dispatch(self, args):
254 '''
255 Dispatch execution to selected command/subcommand.
257 :param args:
258 Parsed arguments obtained from :py:meth:`parse_args`.
260 :returns:
261 ``True`` if dispatching was successful, ``False`` othewise.
263 If an exception of type
264 :py:exc:`~pyrocko.squirrel.error.SquirrelError` or
265 :py:exc:`~pyrocko.squirrel.error.ToolError` is caught, the error is
266 logged and the program is terminated with exit code 1.
267 '''
268 eff_parser = args.__dict__.get('subparser', self)
269 target = args.__dict__.get('target', None)
271 if target:
272 try:
273 target(eff_parser, args)
274 return True
276 except (error.SquirrelError, error.ToolError) as e:
277 logger.fatal(str(e))
278 sys.exit(1)
280 return False
282 def run(self, args=None):
283 '''
284 Parse arguments and dispatch to selected command/subcommand.
286 This simply calls :py:meth:`parse_args` and then :py:meth:`dispatch`
287 with the obtained ``args``. A usage message is printed if no command is
288 selected.
289 '''
290 args = self.parse_args(args)
291 if not self.dispatch(args):
292 self.print_help()
294 def add_squirrel_selection_arguments(self):
295 '''
296 Set up command line options commonly used to configure a
297 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
299 This will optional arguments ``--add``, ``--include``, ``--exclude``,
300 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``, and
301 ``--dataset``.
303 Call ``args.make_squirrel()`` on the arguments returned from
304 :py:meth:`parse_args` to finally instantiate and configure the
305 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
306 '''
307 add_squirrel_selection_arguments(self)
308 self._have_selection_arguments = True
310 def add_squirrel_query_arguments(self, without=[]):
311 '''
312 Set up command line options commonly used in squirrel queries.
314 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
315 ``--tmax``, and ``--time``.
317 Once finished with parsing, the query arguments are available as
318 ``args.squirrel_query`` on the arguments returned from
319 :py:meth:`prase_args`.
321 :param without:
322 Suppress adding given options.
323 :type without:
324 :py:class:`list` of :py:class:`str`, choices: ``'tmin'``,
325 ``'tmax'``, ``'codes'``, and ``'time'``.
326 '''
328 add_squirrel_query_arguments(self, without=without)
329 self._have_query_arguments = True
332def csvtype(choices):
333 def splitarg(arg):
334 values = arg.split(',')
335 for value in values:
336 if value not in choices:
337 raise argparse.ArgumentTypeError(
338 'Invalid choice: {!r} (choose from {})'
339 .format(value, ', '.join(map(repr, choices))))
340 return values
341 return splitarg
344def dq(x):
345 return '``%s``' % x
348def ldq(xs):
349 return ', '.join(dq(x) for x in xs)
352def add_standard_arguments(parser):
353 group = parser.add_argument_group('General options')
354 group.add_argument(
355 '--help', '-h',
356 action='help',
357 help='Show this help message and exit.')
359 loglevel_choices = ['critical', 'error', 'warning', 'info', 'debug']
360 loglevel_default = 'info'
362 group.add_argument(
363 '--loglevel',
364 choices=loglevel_choices,
365 default=loglevel_default,
366 metavar='LEVEL',
367 help='Set logger level. Choices: %s. Default: %s.' % (
368 ldq(loglevel_choices), dq(loglevel_default)))
370 progress_choices = ['terminal', 'log', 'off']
371 progress_default = 'terminal'
373 group.add_argument(
374 '--progress',
375 choices=progress_choices,
376 default=progress_default,
377 metavar='DEST',
378 help='Set how progress status is reported. Choices: %s. '
379 'Default: %s.' % (
380 ldq(progress_choices), dq(progress_default)))
383def process_standard_arguments(parser, args):
384 loglevel = args.__dict__.pop('loglevel')
385 util.setup_logging(parser.prog, loglevel)
387 pmode = args.__dict__.pop('progress')
388 progress.set_default_viewer(pmode)
391def add_squirrel_selection_arguments(parser):
392 '''
393 Set up command line options commonly used to configure a
394 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
396 This will optional arguments ``--add``, ``--include``, ``--exclude``,
397 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``,
398 and ``--dataset`` to a given argument parser.
400 Once finished with parsing, call
401 :py:func:`squirrel_from_selection_arguments` to finally instantiate and
402 configure the :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
404 :param parser:
405 The argument parser to be configured.
406 :type parser:
407 argparse.ArgumentParser
408 '''
409 from pyrocko import squirrel as sq
411 group = parser.add_argument_group('Data collection options')
413 group.add_argument(
414 '--add', '-a',
415 dest='paths',
416 metavar='PATH',
417 nargs='+',
418 help='Add files and directories with waveforms, metadata and events. '
419 'Content is indexed and added to the temporary (default) or '
420 'persistent (see ``--persistent``) data selection.')
422 group.add_argument(
423 '--include',
424 dest='include',
425 metavar='REGEX',
426 help='Only include files whose paths match the regular expression '
427 '``REGEX``. Examples: ``--include=\'\\.MSEED$\'`` would only '
428 'match files ending with ```.MSEED```. '
429 '``--include=\'\\.BH[EN]\\.\'`` would match paths containing '
430 '```.BHE.``` or ```.BHN.```. ``--include=\'/2011/\'`` would '
431 'match paths with a subdirectory ```2011``` in their path '
432 'hierarchy.')
434 group.add_argument(
435 '--exclude',
436 dest='exclude',
437 metavar='REGEX',
438 help='Only include files whose paths do not match the regular '
439 'expression ``REGEX``. Examples: ``--exclude=\'/\\.DS_Store/\'`` '
440 'would exclude anything inside any ```.DS_Store``` subdirectory.')
442 group.add_argument(
443 '--optimistic', '-o',
444 action='store_false',
445 dest='check',
446 default=True,
447 help='Disable checking file modification times for faster startup.')
449 group.add_argument(
450 '--format', '-f',
451 dest='format',
452 metavar='FORMAT',
453 default='detect',
454 choices=sq.supported_formats(),
455 help='Assume input files are of given ``FORMAT``. Choices: %s. '
456 'Default: %s.' % (
457 ldq(sq.supported_formats()),
458 dq('detect')))
460 group.add_argument(
461 '--add-only',
462 type=csvtype(sq.supported_content_kinds()),
463 dest='kinds_add',
464 metavar='KINDS',
465 help='Restrict meta-data scanning to given content kinds. '
466 '``KINDS`` is a comma-separated list of content kinds. '
467 'Choices: %s. By default, all content kinds are indexed.'
468 % ldq(sq.supported_content_kinds()))
470 group.add_argument(
471 '--persistent', '-p',
472 dest='persistent',
473 metavar='NAME',
474 help='Create/use persistent selection with given ``NAME``. Persistent '
475 'selections can be used to speed up startup of Squirrel-based '
476 'applications.')
478 group.add_argument(
479 '--dataset', '-d',
480 dest='datasets',
481 default=[],
482 action='append',
483 metavar='FILE',
484 help='Add files, directories and remote sources from dataset '
485 'description file. This option can be repeated to add multiple '
486 'datasets. Run ```squirrel template``` to obtain examples of '
487 'dataset description files.')
490def squirrel_from_selection_arguments(args):
491 '''
492 Create a :py:class:`~pyrocko.squirrel.base.Squirrel` instance from command
493 line arguments.
495 Use :py:func:`add_squirrel_selection_arguments` to configure the parser
496 with the necessary options.
498 :param args:
499 Parsed command line arguments, as returned by
500 :py:meth:`argparse.ArgumentParser.parse_args`.
502 :returns:
503 :py:class:`pyrocko.squirrel.base.Squirrel` instance with paths,
504 datasets and remote sources added.
506 '''
507 from pyrocko.squirrel import base
509 squirrel = base.Squirrel(persistent=args.persistent)
511 if args.paths:
512 squirrel.add(
513 args.paths,
514 check=args.check,
515 format=args.format,
516 kinds=args.kinds_add or None,
517 include=args.include,
518 exclude=args.exclude)
520 for dataset_path in args.datasets:
521 squirrel.add_dataset(dataset_path, check=args.check)
523 return squirrel
526def add_squirrel_query_arguments(parser, without=[]):
527 '''
528 Set up command line options commonly used in squirrel queries.
530 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
531 ``--tmax``, and ``--time``.
533 Once finished with parsing, call
534 :py:func:`squirrel_query_from_arguments` to get the parsed values.
536 :param parser:
537 The argument parser to be configured.
538 :type parser:
539 argparse.ArgumentParser
541 :param without:
542 Suppress adding given options.
543 :type without:
544 :py:class:`list` of :py:class:`str`
545 '''
547 from pyrocko import squirrel as sq
549 group = parser.add_argument_group('Data query options')
551 if 'kinds' not in without:
552 group.add_argument(
553 '--kinds',
554 type=csvtype(sq.supported_content_kinds()),
555 dest='kinds',
556 metavar='KINDS',
557 help='Content kinds to query. ``KINDS`` is a comma-separated list '
558 'of content kinds. Choices: %s. By default, all content '
559 'kinds are queried.' % ldq(sq.supported_content_kinds()))
561 if 'codes' not in without:
562 group.add_argument(
563 '--codes',
564 dest='codes',
565 metavar='CODES',
566 help='Code patterns to query (``STA``, ``NET.STA``, '
567 '``NET.STA.LOC``, ``NET.STA.LOC.CHA``, or '
568 '``NET.STA.LOC.CHA.EXTRA``). The pattern may contain '
569 'wildcards ``*`` (zero or more arbitrary characters), ``?`` '
570 '(single arbitrary character), and ``[CHARS]`` (any '
571 'character out of ``CHARS``). Multiple patterns can be '
572 'given by separating them with commas.')
574 if 'tmin' not in without:
575 group.add_argument(
576 '--tmin',
577 dest='tmin',
578 metavar='TIME',
579 help='Begin of time interval to query. %s' % help_time_format)
581 if 'tmax' not in without:
582 group.add_argument(
583 '--tmax',
584 dest='tmax',
585 metavar='TIME',
586 help='End of time interval to query. %s' % help_time_format)
588 if 'time' not in without:
589 group.add_argument(
590 '--time',
591 dest='time',
592 metavar='TIME',
593 help='Time instant to query. %s' % help_time_format)
596def squirrel_query_from_arguments(args):
597 '''
598 Get common arguments to be used in squirrel queries from command line.
600 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
601 the necessary options.
603 :param args:
604 Parsed command line arguments, as returned by
605 :py:meth:`argparse.ArgumentParser.parse_args`.
607 :returns:
608 :py:class:`dict` with any parsed option values.
609 '''
611 from pyrocko import squirrel as sq
613 d = {}
615 if 'kinds' in args and args.kinds:
616 d['kind'] = args.kinds
617 if 'tmin' in args and args.tmin:
618 d['tmin'] = util.str_to_time_fillup(args.tmin)
619 if 'tmax' in args and args.tmax:
620 d['tmax'] = util.str_to_time_fillup(args.tmax)
621 if 'time' in args and args.time:
622 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
623 if 'codes' in args and args.codes:
624 d['codes'] = [
625 sq.to_codes_guess(s.strip()) for s in args.codes.split(',')]
627 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
628 raise error.SquirrelError(
629 'Options --tmin/--tmax and --time are mutually exclusive.')
630 return d
633class SquirrelCommand(object):
634 '''
635 Base class for Squirrel-based CLI programs and subcommands.
636 '''
638 def fail(self, message):
639 '''
640 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
642 :py:func:`~pyrocko.squirrel.tool.from_command` catches
643 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
644 terminates with an error exit state.
645 '''
646 raise error.ToolError(message)
648 def make_subparser(self, subparsers):
649 '''
650 To be implemented in subcommand. Create subcommand parser.
652 Must return a newly created parser obtained with
653 :py:meth:`add_parser`, e.g.::
655 def make_subparser(self, subparsers):
656 return subparsers.add_parser(
657 'plot', help='Draw a nice plot.')
659 '''
660 return subparsers.add_parser(
661 self.__class__.__name__, help='Undocumented.')
663 def setup(self, parser):
664 '''
665 To be implemented in subcommand. Configure parser.
667 :param parser:
668 The argument parser to be configured.
669 :type parser:
670 argparse.ArgumentParser
672 Example::
674 def setup(self, parser):
675 parser.add_squirrel_selection_arguments()
676 parser.add_squirrel_query_arguments()
677 parser.add_argument(
678 '--fmin',
679 dest='fmin',
680 metavar='FLOAT',
681 type=float,
682 help='Corner of highpass [Hz].')
683 '''
684 pass
686 def run(self, parser, args):
687 '''
688 To be implemented in subcommand. Main routine of the command.
690 :param parser:
691 The argument parser to be configured.
692 :type parser:
693 argparse.ArgumentParser
695 :param args:
696 Parsed command line arguments, as returned by
697 :py:meth:`argparse.ArgumentParser.parse_args`.
699 Example::
701 def run(self, parser, args):
702 print('User has selected fmin = %g Hz' % args.fmin)
704 # args.make_squirrel() is available if
705 # parser.add_squirrel_selection_arguments() was called during
706 # setup().
708 sq = args.make_squirrel()
710 # args.squirrel_query is available if
711 # praser.add_squirrel_query_arguments() was called during
712 # setup().
714 stations = sq.get_stations(**args.squirrel_query)
715 '''
716 pass
719__all__ = [
720 'SquirrelArgumentParser',
721 'SquirrelCommand',
722 'add_squirrel_selection_arguments',
723 'squirrel_from_selection_arguments',
724 'add_squirrel_query_arguments',
725 'squirrel_query_from_arguments',
726]