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``,
301 and ``--update``, and ``--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 ``--update``, 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 '--update', '-u',
480 dest='update',
481 action='store_true',
482 default=False,
483 help='Allow adding paths and datasets to existing persistent '
484 'selection.')
486 group.add_argument(
487 '--dataset', '-d',
488 dest='datasets',
489 default=[],
490 action='append',
491 metavar='FILE',
492 help='Add files, directories and remote sources from dataset '
493 'description file. This option can be repeated to add multiple '
494 'datasets. Run ```squirrel template``` to obtain examples of '
495 'dataset description files.')
498def squirrel_from_selection_arguments(args):
499 '''
500 Create a :py:class:`~pyrocko.squirrel.base.Squirrel` instance from command
501 line arguments.
503 Use :py:func:`add_squirrel_selection_arguments` to configure the parser
504 with the necessary options.
506 :param args:
507 Parsed command line arguments, as returned by
508 :py:meth:`argparse.ArgumentParser.parse_args`.
510 :returns:
511 :py:class:`pyrocko.squirrel.base.Squirrel` instance with paths,
512 datasets and remote sources added.
514 '''
515 from pyrocko.squirrel import base, dataset
517 datasets = [
518 dataset.read_dataset(dataset_path) for dataset_path in args.datasets]
520 persistents = [ds.persistent or '' for ds in datasets if ds.persistent]
521 if args.persistent:
522 persistent = args.persistent
523 elif persistents:
524 persistent = persistents[0]
525 if not all(p == persistents for p in persistents[1:]):
526 raise error.SquirrelError(
527 'Given datasets specify different `persistent` settings.')
529 if persistent:
530 logger.info(
531 'Persistent selection requested by dataset: %s' % persistent)
532 else:
533 persistent = None
535 else:
536 persistent = None
538 squirrel = base.Squirrel(persistent=persistent)
540 if persistent and not squirrel.is_new():
541 if not args.update:
542 logger.info(
543 'Using existing persistent selection: %s' % persistent)
544 if args.paths or datasets:
545 logger.info(
546 'Avoiding dataset rescan. Use --update/-u to '
547 'rescan or add items to existing persistent selection.')
549 return squirrel
551 else:
552 logger.info(
553 'Updating existing persistent selection: %s' % persistent)
555 if args.paths:
556 squirrel.add(
557 args.paths,
558 check=args.check,
559 format=args.format,
560 kinds=args.kinds_add or None,
561 include=args.include,
562 exclude=args.exclude)
564 for ds in datasets:
565 squirrel.add_dataset(ds, check=args.check)
567 return squirrel
570def add_squirrel_query_arguments(parser, without=[]):
571 '''
572 Set up command line options commonly used in squirrel queries.
574 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
575 ``--tmax``, and ``--time``.
577 Once finished with parsing, call
578 :py:func:`squirrel_query_from_arguments` to get the parsed values.
580 :param parser:
581 The argument parser to be configured.
582 :type parser:
583 argparse.ArgumentParser
585 :param without:
586 Suppress adding given options.
587 :type without:
588 :py:class:`list` of :py:class:`str`
589 '''
591 from pyrocko import squirrel as sq
593 group = parser.add_argument_group('Data query options')
595 if 'kinds' not in without:
596 group.add_argument(
597 '--kinds',
598 type=csvtype(sq.supported_content_kinds()),
599 dest='kinds',
600 metavar='KINDS',
601 help='Content kinds to query. ``KINDS`` is a comma-separated list '
602 'of content kinds. Choices: %s. By default, all content '
603 'kinds are queried.' % ldq(sq.supported_content_kinds()))
605 if 'codes' not in without:
606 group.add_argument(
607 '--codes',
608 dest='codes',
609 metavar='CODES',
610 help='Code patterns to query (``STA``, ``NET.STA``, '
611 '``NET.STA.LOC``, ``NET.STA.LOC.CHA``, or '
612 '``NET.STA.LOC.CHA.EXTRA``). The pattern may contain '
613 'wildcards ``*`` (zero or more arbitrary characters), ``?`` '
614 '(single arbitrary character), and ``[CHARS]`` (any '
615 'character out of ``CHARS``). Multiple patterns can be '
616 'given by separating them with commas.')
618 if 'tmin' not in without:
619 group.add_argument(
620 '--tmin',
621 dest='tmin',
622 metavar='TIME',
623 help='Begin of time interval to query. %s' % help_time_format)
625 if 'tmax' not in without:
626 group.add_argument(
627 '--tmax',
628 dest='tmax',
629 metavar='TIME',
630 help='End of time interval to query. %s' % help_time_format)
632 if 'time' not in without:
633 group.add_argument(
634 '--time',
635 dest='time',
636 metavar='TIME',
637 help='Time instant to query. %s' % help_time_format)
640def squirrel_query_from_arguments(args):
641 '''
642 Get common arguments to be used in squirrel queries from command line.
644 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
645 the necessary options.
647 :param args:
648 Parsed command line arguments, as returned by
649 :py:meth:`argparse.ArgumentParser.parse_args`.
651 :returns:
652 :py:class:`dict` with any parsed option values.
653 '''
655 from pyrocko import squirrel as sq
657 d = {}
659 if 'kinds' in args and args.kinds:
660 d['kind'] = args.kinds
661 if 'tmin' in args and args.tmin:
662 d['tmin'] = util.str_to_time_fillup(args.tmin)
663 if 'tmax' in args and args.tmax:
664 d['tmax'] = util.str_to_time_fillup(args.tmax)
665 if 'time' in args and args.time:
666 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
667 if 'codes' in args and args.codes:
668 d['codes'] = [
669 sq.to_codes_guess(s.strip()) for s in args.codes.split(',')]
671 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
672 raise error.SquirrelError(
673 'Options --tmin/--tmax and --time are mutually exclusive.')
674 return d
677class SquirrelCommand(object):
678 '''
679 Base class for Squirrel-based CLI programs and subcommands.
680 '''
682 def fail(self, message):
683 '''
684 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
686 :py:func:`~pyrocko.squirrel.tool.from_command` catches
687 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
688 terminates with an error exit state.
689 '''
690 raise error.ToolError(message)
692 def make_subparser(self, subparsers):
693 '''
694 To be implemented in subcommand. Create subcommand parser.
696 Must return a newly created parser obtained with
697 :py:meth:`add_parser`, e.g.::
699 def make_subparser(self, subparsers):
700 return subparsers.add_parser(
701 'plot', help='Draw a nice plot.')
703 '''
704 return subparsers.add_parser(
705 self.__class__.__name__, help='Undocumented.')
707 def setup(self, parser):
708 '''
709 To be implemented in subcommand. Configure parser.
711 :param parser:
712 The argument parser to be configured.
713 :type parser:
714 argparse.ArgumentParser
716 Example::
718 def setup(self, parser):
719 parser.add_squirrel_selection_arguments()
720 parser.add_squirrel_query_arguments()
721 parser.add_argument(
722 '--fmin',
723 dest='fmin',
724 metavar='FLOAT',
725 type=float,
726 help='Corner of highpass [Hz].')
727 '''
728 pass
730 def run(self, parser, args):
731 '''
732 To be implemented in subcommand. Main routine of the command.
734 :param parser:
735 The argument parser to be configured.
736 :type parser:
737 argparse.ArgumentParser
739 :param args:
740 Parsed command line arguments, as returned by
741 :py:meth:`argparse.ArgumentParser.parse_args`.
743 Example::
745 def run(self, parser, args):
746 print('User has selected fmin = %g Hz' % args.fmin)
748 # args.make_squirrel() is available if
749 # parser.add_squirrel_selection_arguments() was called during
750 # setup().
752 sq = args.make_squirrel()
754 # args.squirrel_query is available if
755 # praser.add_squirrel_query_arguments() was called during
756 # setup().
758 stations = sq.get_stations(**args.squirrel_query)
759 '''
760 pass
763__all__ = [
764 'SquirrelArgumentParser',
765 'SquirrelCommand',
766 'add_squirrel_selection_arguments',
767 'squirrel_from_selection_arguments',
768 'add_squirrel_query_arguments',
769 'squirrel_query_from_arguments',
770]