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 if part.startswith('usage:'):
47 lines.extend(part.splitlines())
48 else:
49 for line in part.splitlines():
50 if not line:
51 lines.append(line)
52 if not line.startswith(' '):
53 lines.extend(
54 textwrap.wrap(line, 79,))
55 else:
56 lines.extend(
57 textwrap.wrap(line, 79, subsequent_indent=' '*24))
59 lines.append('')
61 return '\n'.join(lines)
64def wrap_usage(s):
65 lines = []
66 for line in s.splitlines():
67 if not line.startswith('usage:'):
68 lines.append(line)
69 else:
70 lines.extend(textwrap.wrap(line, 79, subsequent_indent=' '*24))
72 return '\n'.join(lines)
75def formatter_with_width(n):
76 class PyrockoHelpFormatter(argparse.RawDescriptionHelpFormatter):
77 def __init__(self, *args, **kwargs):
78 kwargs['width'] = n
79 kwargs['max_help_position'] = 24
80 argparse.RawDescriptionHelpFormatter.__init__(
81 self, *args, **kwargs)
83 # fix alignment problems, with the post-processing wrapping
84 self._action_max_length = 24
86 return PyrockoHelpFormatter
89class PyrockoArgumentParser(argparse.ArgumentParser):
91 # We want to convert the --help outputs to rst for the html docs. Problem
92 # is that argparse's HelpFormatters to date have no public interface which
93 # we could use to achieve this. The solution here is a bit clunky but works
94 # ok for Squirrel. We allow markup like ``code`` which is kept when
95 # producing rst (by parsing the final --help output) but stripped out when
96 # doing normal --help. This leads to a problem with the internal wrapping
97 # of argparse does this before the stripping. To solve, we render with
98 # argparse to a very wide width and do the wrapping in postprocessing.
99 # ``code`` is replaced with just code in normal output. ```code``` is
100 # replaced with 'code' in normal output and with ``code`` rst output. rst
101 # output is selected with environment variable PYROCKO_RST_HELP=1.
102 # The script maintenance/argparse_help_to_rst.py extracts the rst help
103 # and generates the rst files for the docs.
105 def __init__(
106 self, prog=None, usage=None, description=None, epilog=None,
107 **kwargs):
109 kwargs['formatter_class'] = formatter_with_width(1000000)
111 description = unwrap(description)
112 epilog = unwrap(epilog)
114 argparse.ArgumentParser.__init__(
115 self, prog=prog, usage=usage, description=description,
116 epilog=epilog, **kwargs)
118 if hasattr(self, '_action_groups'):
119 for group in self._action_groups:
120 if group.title == 'positional arguments':
121 group.title = 'Positional arguments'
123 elif group.title == 'optional arguments':
124 group.title = 'Optional arguments'
126 elif group.title == 'options':
127 group.title = 'Options'
129 self.raw_help = False
131 def format_help(self, *args, **kwargs):
132 s = argparse.ArgumentParser.format_help(self, *args, **kwargs)
134 # replace usage with wrapped one from argparse because naive wrapping
135 # does not look good.
136 formatter_class = self.formatter_class
137 self.formatter_class = formatter_with_width(79)
138 usage = self.format_usage()
139 self.formatter_class = formatter_class
141 lines = []
142 for line in s.splitlines():
143 if line.startswith('usage:'):
144 lines.append(usage)
145 else:
146 lines.append(line)
148 s = '\n'.join(lines)
150 if os.environ.get('PYROCKO_RST_HELP', '0') == '0':
151 s = s.replace('```', '\'')
152 s = s.replace('``', '')
153 s = wrap(s)
154 else:
155 s = s.replace('```', '``')
156 s = wrap_usage(s)
158 return s
161class SquirrelArgumentParser(PyrockoArgumentParser):
162 '''
163 Parser for CLI arguments with a some extras for Squirrel based apps.
165 :param command:
166 Implementation of the command.
167 :type command:
168 :py:class:`SquirrelCommand` (or module providing the same interface).
170 :param subcommands:
171 Implementations of subcommands.
172 :type subcommands:
173 :py:class:`list` of :py:class:`SquirrelCommand` (or modules providing
174 the same interface).
176 :param \\*args:
177 Handed through to base class's init.
179 :param \\*\\*kwargs:
180 Handed through to base class's init.
181 '''
183 def __init__(self, *args, command=None, subcommands=[], **kwargs):
185 self._command = command
186 self._subcommands = subcommands
187 self._have_selection_arguments = False
188 self._have_query_arguments = False
190 kwargs['add_help'] = False
191 PyrockoArgumentParser.__init__(self, *args, **kwargs)
192 add_standard_arguments(self)
193 self._command = None
194 self._subcommands = []
195 if command:
196 self.set_command(command)
198 if subcommands:
199 self.set_subcommands(subcommands)
201 def set_command(self, command):
202 command.setup(self)
203 self.set_defaults(target=command.run)
205 def set_subcommands(self, subcommands):
206 subparsers = self.add_subparsers(
207 metavar='SUBCOMMAND',
208 title='Subcommands')
210 for mod in subcommands:
211 subparser = mod.make_subparser(subparsers)
212 if subparser is None:
213 raise Exception(
214 'make_subparser(subparsers) must return the created '
215 'parser.')
217 mod.setup(subparser)
218 subparser.set_defaults(target=mod.run, subparser=subparser)
220 def parse_args(self, args=None, namespace=None):
221 '''
222 Parse arguments given on command line.
224 Extends the functionality of
225 :py:meth:`argparse.ArgumentParser.parse_args` to process and handle the
226 standard options ``--loglevel``, ``--progress`` and ``--help``.
227 '''
229 args = PyrockoArgumentParser.parse_args(
230 self, args=args, namespace=namespace)
232 eff_parser = args.__dict__.get('subparser', self)
234 process_standard_arguments(eff_parser, args)
236 if eff_parser._have_selection_arguments:
237 def make_squirrel():
238 return squirrel_from_selection_arguments(args)
240 args.make_squirrel = make_squirrel
242 if eff_parser._have_query_arguments:
243 try:
244 args.squirrel_query = squirrel_query_from_arguments(args)
245 except (error.SquirrelError, error.ToolError) as e:
246 logger.fatal(str(e))
247 sys.exit(1)
249 return args
251 def dispatch(self, args):
252 '''
253 Dispatch execution to selected command/subcommand.
255 :param args:
256 Parsed arguments obtained from :py:meth:`parse_args`.
258 :returns:
259 ``True`` if dispatching was successful, ``False`` othewise.
261 If an exception of type
262 :py:exc:`~pyrocko.squirrel.error.SquirrelError` or
263 :py:exc:`~pyrocko.squirrel.error.ToolError` is caught, the error is
264 logged and the program is terminated with exit code 1.
265 '''
266 eff_parser = args.__dict__.get('subparser', self)
267 target = args.__dict__.get('target', None)
269 if target:
270 try:
271 target(eff_parser, args)
272 return True
274 except (error.SquirrelError, error.ToolError) as e:
275 logger.fatal(str(e))
276 sys.exit(1)
278 return False
280 def run(self, args=None):
281 '''
282 Parse arguments and dispatch to selected command/subcommand.
284 This simply calls :py:meth:`parse_args` and then :py:meth:`dispatch`
285 with the obtained ``args``. A usage message is printed if no command is
286 selected.
287 '''
288 args = self.parse_args(args)
289 if not self.dispatch(args):
290 self.print_help()
292 def add_squirrel_selection_arguments(self):
293 '''
294 Set up command line options commonly used to configure a
295 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
297 This will optional arguments ``--add``, ``--include``, ``--exclude``,
298 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``,
299 and ``--update``, and ``--dataset``.
301 Call ``args.make_squirrel()`` on the arguments returned from
302 :py:meth:`parse_args` to finally instantiate and configure the
303 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
304 '''
305 add_squirrel_selection_arguments(self)
306 self._have_selection_arguments = True
308 def add_squirrel_query_arguments(self, without=[]):
309 '''
310 Set up command line options commonly used in squirrel queries.
312 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
313 ``--tmax``, and ``--time``.
315 Once finished with parsing, the query arguments are available as
316 ``args.squirrel_query`` on the arguments returned from
317 :py:meth:`prase_args`.
319 :param without:
320 Suppress adding given options.
321 :type without:
322 :py:class:`list` of :py:class:`str`, choices: ``'tmin'``,
323 ``'tmax'``, ``'codes'``, and ``'time'``.
324 '''
326 add_squirrel_query_arguments(self, without=without)
327 self._have_query_arguments = True
330def csvtype(choices):
331 def splitarg(arg):
332 values = arg.split(',')
333 for value in values:
334 if value not in choices:
335 raise argparse.ArgumentTypeError(
336 'Invalid choice: {!r} (choose from {})'
337 .format(value, ', '.join(map(repr, choices))))
338 return values
339 return splitarg
342def dq(x):
343 return '``%s``' % x
346def ldq(xs):
347 return ', '.join(dq(x) for x in xs)
350def add_standard_arguments(parser):
351 group = parser.add_argument_group('General options')
352 group.add_argument(
353 '--help', '-h',
354 action='help',
355 help='Show this help message and exit.')
357 loglevel_choices = ['critical', 'error', 'warning', 'info', 'debug']
358 loglevel_default = 'info'
360 group.add_argument(
361 '--loglevel',
362 choices=loglevel_choices,
363 default=loglevel_default,
364 metavar='LEVEL',
365 help='Set logger level. Choices: %s. Default: %s.' % (
366 ldq(loglevel_choices), dq(loglevel_default)))
368 progress_choices = ['terminal', 'log', 'off']
369 progress_default = 'terminal'
371 group.add_argument(
372 '--progress',
373 choices=progress_choices,
374 default=progress_default,
375 metavar='DEST',
376 help='Set how progress status is reported. Choices: %s. '
377 'Default: %s.' % (
378 ldq(progress_choices), dq(progress_default)))
381def process_standard_arguments(parser, args):
382 loglevel = args.__dict__.pop('loglevel')
383 util.setup_logging(parser.prog, loglevel)
385 pmode = args.__dict__.pop('progress')
386 progress.set_default_viewer(pmode)
389def add_squirrel_selection_arguments(parser):
390 '''
391 Set up command line options commonly used to configure a
392 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
394 This will optional arguments ``--add``, ``--include``, ``--exclude``,
395 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``,
396 and ``--update``, and ``--dataset`` to a given argument parser.
398 Once finished with parsing, call
399 :py:func:`squirrel_from_selection_arguments` to finally instantiate and
400 configure the :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
402 :param parser:
403 The argument parser to be configured.
404 :type parser:
405 argparse.ArgumentParser
406 '''
407 from pyrocko import squirrel as sq
409 group = parser.add_argument_group('Data collection options')
411 group.add_argument(
412 '--add', '-a',
413 dest='paths',
414 metavar='PATH',
415 nargs='+',
416 help='Add files and directories with waveforms, metadata and events. '
417 'Content is indexed and added to the temporary (default) or '
418 'persistent (see ``--persistent``) data selection.')
420 group.add_argument(
421 '--include',
422 dest='include',
423 metavar='REGEX',
424 help='Only include files whose paths match the regular expression '
425 '``REGEX``. Examples: ``--include=\'\\.MSEED$\'`` would only '
426 'match files ending with ```.MSEED```. '
427 '``--include=\'\\.BH[EN]\\.\'`` would match paths containing '
428 '```.BHE.``` or ```.BHN.```. ``--include=\'/2011/\'`` would '
429 'match paths with a subdirectory ```2011``` in their path '
430 'hierarchy.')
432 group.add_argument(
433 '--exclude',
434 dest='exclude',
435 metavar='REGEX',
436 help='Only include files whose paths do not match the regular '
437 'expression ``REGEX``. Examples: ``--exclude=\'/\\.DS_Store/\'`` '
438 'would exclude anything inside any ```.DS_Store``` subdirectory.')
440 group.add_argument(
441 '--optimistic', '-o',
442 action='store_false',
443 dest='check',
444 default=True,
445 help='Disable checking file modification times for faster startup.')
447 group.add_argument(
448 '--format', '-f',
449 dest='format',
450 metavar='FORMAT',
451 default='detect',
452 choices=sq.supported_formats(),
453 help='Assume input files are of given ``FORMAT``. Choices: %s. '
454 'Default: %s.' % (
455 ldq(sq.supported_formats()),
456 dq('detect')))
458 group.add_argument(
459 '--add-only',
460 type=csvtype(sq.supported_content_kinds()),
461 dest='kinds_add',
462 metavar='KINDS',
463 help='Restrict meta-data scanning to given content kinds. '
464 '``KINDS`` is a comma-separated list of content kinds. '
465 'Choices: %s. By default, all content kinds are indexed.'
466 % ldq(sq.supported_content_kinds()))
468 group.add_argument(
469 '--persistent', '-p',
470 dest='persistent',
471 metavar='NAME',
472 help='Create/use persistent selection with given ``NAME``. Persistent '
473 'selections can be used to speed up startup of Squirrel-based '
474 'applications.')
476 group.add_argument(
477 '--update', '-u',
478 dest='update',
479 action='store_true',
480 default=False,
481 help='Allow adding paths and datasets to existing persistent '
482 'selection.')
484 group.add_argument(
485 '--dataset', '-d',
486 dest='datasets',
487 default=[],
488 action='append',
489 metavar='FILE',
490 help='Add files, directories and remote sources from dataset '
491 'description file. This option can be repeated to add multiple '
492 'datasets. Run ```squirrel template``` to obtain examples of '
493 'dataset description files.')
496def squirrel_from_selection_arguments(args):
497 '''
498 Create a :py:class:`~pyrocko.squirrel.base.Squirrel` instance from command
499 line arguments.
501 Use :py:func:`add_squirrel_selection_arguments` to configure the parser
502 with the necessary options.
504 :param args:
505 Parsed command line arguments, as returned by
506 :py:meth:`argparse.ArgumentParser.parse_args`.
508 :returns:
509 :py:class:`pyrocko.squirrel.base.Squirrel` instance with paths,
510 datasets and remote sources added.
512 '''
513 from pyrocko.squirrel import base, dataset
515 datasets = [
516 dataset.read_dataset(dataset_path) for dataset_path in args.datasets]
518 persistents = [ds.persistent or '' for ds in datasets if ds.persistent]
519 if args.persistent:
520 persistent = args.persistent
521 elif persistents:
522 persistent = persistents[0]
523 if not all(p == persistents for p in persistents[1:]):
524 raise error.SquirrelError(
525 'Given datasets specify different `persistent` settings.')
527 if persistent:
528 logger.info(
529 'Persistent selection requested by dataset: %s' % persistent)
530 else:
531 persistent = None
533 else:
534 persistent = None
536 squirrel = base.Squirrel(persistent=persistent)
538 if persistent and not squirrel.is_new():
539 if not args.update:
540 logger.info(
541 'Using existing persistent selection: %s' % persistent)
542 if args.paths or datasets:
543 logger.info(
544 'Avoiding dataset rescan. Use --update/-u to '
545 'rescan or add items to existing persistent selection.')
547 return squirrel
549 else:
550 logger.info(
551 'Updating existing persistent selection: %s' % persistent)
553 if args.paths:
554 squirrel.add(
555 args.paths,
556 check=args.check,
557 format=args.format,
558 kinds=args.kinds_add or None,
559 include=args.include,
560 exclude=args.exclude)
562 for ds in datasets:
563 squirrel.add_dataset(ds, check=args.check)
565 return squirrel
568def add_squirrel_query_arguments(parser, without=[]):
569 '''
570 Set up command line options commonly used in squirrel queries.
572 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
573 ``--tmax``, and ``--time``.
575 Once finished with parsing, call
576 :py:func:`squirrel_query_from_arguments` to get the parsed values.
578 :param parser:
579 The argument parser to be configured.
580 :type parser:
581 argparse.ArgumentParser
583 :param without:
584 Suppress adding given options.
585 :type without:
586 :py:class:`list` of :py:class:`str`
587 '''
589 from pyrocko import squirrel as sq
591 group = parser.add_argument_group('Data query options')
593 if 'kinds' not in without:
594 group.add_argument(
595 '--kinds',
596 type=csvtype(sq.supported_content_kinds()),
597 dest='kinds',
598 metavar='KINDS',
599 help='Content kinds to query. ``KINDS`` is a comma-separated list '
600 'of content kinds. Choices: %s. By default, all content '
601 'kinds are queried.' % ldq(sq.supported_content_kinds()))
603 if 'codes' not in without:
604 group.add_argument(
605 '--codes',
606 dest='codes',
607 metavar='CODES',
608 help='Code patterns to query (``STA``, ``NET.STA``, '
609 '``NET.STA.LOC``, ``NET.STA.LOC.CHA``, or '
610 '``NET.STA.LOC.CHA.EXTRA``). The pattern may contain '
611 'wildcards ``*`` (zero or more arbitrary characters), ``?`` '
612 '(single arbitrary character), and ``[CHARS]`` (any '
613 'character out of ``CHARS``). Multiple patterns can be '
614 'given by separating them with commas.')
616 if 'tmin' not in without:
617 group.add_argument(
618 '--tmin',
619 dest='tmin',
620 metavar='TIME',
621 help='Begin of time interval to query. %s' % help_time_format)
623 if 'tmax' not in without:
624 group.add_argument(
625 '--tmax',
626 dest='tmax',
627 metavar='TIME',
628 help='End of time interval to query. %s' % help_time_format)
630 if 'time' not in without:
631 group.add_argument(
632 '--time',
633 dest='time',
634 metavar='TIME',
635 help='Time instant to query. %s' % help_time_format)
638def squirrel_query_from_arguments(args):
639 '''
640 Get common arguments to be used in squirrel queries from command line.
642 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
643 the necessary options.
645 :param args:
646 Parsed command line arguments, as returned by
647 :py:meth:`argparse.ArgumentParser.parse_args`.
649 :returns:
650 :py:class:`dict` with any parsed option values.
651 '''
653 from pyrocko import squirrel as sq
655 d = {}
657 if 'kinds' in args and args.kinds:
658 d['kind'] = args.kinds
659 if 'tmin' in args and args.tmin:
660 d['tmin'] = util.str_to_time_fillup(args.tmin)
661 if 'tmax' in args and args.tmax:
662 d['tmax'] = util.str_to_time_fillup(args.tmax)
663 if 'time' in args and args.time:
664 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
665 if 'codes' in args and args.codes:
666 d['codes'] = [
667 sq.to_codes_guess(s.strip()) for s in args.codes.split(',')]
669 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
670 raise error.SquirrelError(
671 'Options --tmin/--tmax and --time are mutually exclusive.')
672 return d
675class SquirrelCommand(object):
676 '''
677 Base class for Squirrel-based CLI programs and subcommands.
678 '''
680 def fail(self, message):
681 '''
682 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
684 :py:func:`~pyrocko.squirrel.tool.from_command` catches
685 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
686 terminates with an error exit state.
687 '''
688 raise error.ToolError(message)
690 def make_subparser(self, subparsers):
691 '''
692 To be implemented in subcommand. Create subcommand parser.
694 Must return a newly created parser obtained with
695 :py:meth:`add_parser`, e.g.::
697 def make_subparser(self, subparsers):
698 return subparsers.add_parser(
699 'plot', help='Draw a nice plot.')
701 '''
702 return subparsers.add_parser(
703 self.__class__.__name__, help='Undocumented.')
705 def setup(self, parser):
706 '''
707 To be implemented in subcommand. Configure parser.
709 :param parser:
710 The argument parser to be configured.
711 :type parser:
712 argparse.ArgumentParser
714 Example::
716 def setup(self, parser):
717 parser.add_squirrel_selection_arguments()
718 parser.add_squirrel_query_arguments()
719 parser.add_argument(
720 '--fmin',
721 dest='fmin',
722 metavar='FLOAT',
723 type=float,
724 help='Corner of highpass [Hz].')
725 '''
726 pass
728 def run(self, parser, args):
729 '''
730 To be implemented in subcommand. Main routine of the command.
732 :param parser:
733 The argument parser to be configured.
734 :type parser:
735 argparse.ArgumentParser
737 :param args:
738 Parsed command line arguments, as returned by
739 :py:meth:`argparse.ArgumentParser.parse_args`.
741 Example::
743 def run(self, parser, args):
744 print('User has selected fmin = %g Hz' % args.fmin)
746 # args.make_squirrel() is available if
747 # parser.add_squirrel_selection_arguments() was called during
748 # setup().
750 sq = args.make_squirrel()
752 # args.squirrel_query is available if
753 # praser.add_squirrel_query_arguments() was called during
754 # setup().
756 stations = sq.get_stations(**args.squirrel_query)
757 '''
758 pass
761__all__ = [
762 'SquirrelArgumentParser',
763 'SquirrelCommand',
764 'add_squirrel_selection_arguments',
765 'squirrel_from_selection_arguments',
766 'add_squirrel_query_arguments',
767 'squirrel_query_from_arguments',
768]