Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/squirrel/tool/common.py: 79%
230 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-10-06 15:01 +0000
1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6'''
7Squirrel command line tool infrastructure and argument parsing.
8'''
10import os
11import sys
12import re
13import argparse
14import logging
15import textwrap
17from pyrocko import util, progress
18from pyrocko.squirrel import error
21logger = logging.getLogger('psq.tool.common')
23help_time_format = 'Format: ```YYYY-MM-DD HH:MM:SS.FFF```, truncation allowed.'
26def unwrap(s):
27 if s is None:
28 return None
29 s = s.strip()
30 parts = re.split(r'\n{2,}', s)
31 lines = []
32 for part in parts:
33 plines = part.splitlines()
34 if not any(line.startswith(' ') for line in plines):
35 lines.append(' '.join(plines))
36 else:
37 lines.extend(plines)
39 lines.append('')
41 return '\n'.join(lines)
44def wrap(s):
45 lines = []
46 parts = re.split(r'\n{2,}', s)
47 for part in parts:
48 plines = part.splitlines()
49 if part.startswith('usage:') \
50 or all(line.startswith(' ') for line in plines):
51 lines.extend(plines)
52 else:
53 for line in plines:
54 if not line:
55 lines.append(line)
56 if not line.startswith(' '):
57 lines.extend(
58 textwrap.wrap(line, 79,))
59 else:
60 lines.extend(
61 textwrap.wrap(line, 79, subsequent_indent=' '*24))
63 lines.append('')
65 return '\n'.join(lines)
68def wrap_usage(s):
69 lines = []
70 for line in s.splitlines():
71 if not line.startswith('usage:'):
72 lines.append(line)
73 else:
74 lines.extend(textwrap.wrap(line, 79, subsequent_indent=' '*24))
76 return '\n'.join(lines)
79def formatter_with_width(n):
80 class PyrockoHelpFormatter(argparse.RawDescriptionHelpFormatter):
81 def __init__(self, *args, **kwargs):
82 kwargs['width'] = n
83 kwargs['max_help_position'] = 24
84 argparse.RawDescriptionHelpFormatter.__init__(
85 self, *args, **kwargs)
87 # fix alignment problems, with the post-processing wrapping
88 self._action_max_length = 24
90 return PyrockoHelpFormatter
93class PyrockoArgumentParser(argparse.ArgumentParser):
94 '''
95 Tweaks and extends the standard argument parser to simplify the generation
96 of the online docs.
98 We want to convert the ``--help`` outputs to ``rst`` for the html docs.
99 Problem is that argparse's ``HelpFormatter`` to date have no public
100 interface which we could use to achieve this. The solution here is a bit
101 clunky but works ok for our purposes. We allow markup like
102 :literal:`\\`\\`code\\`\\`` which is kept when producing ``rst`` (by
103 parsing the final ``--help`` output) but stripped out when doing normal
104 ``--help``. This leads to a problem with the internal output wrapping of
105 argparse which it does before the stripping. To solve, we render with
106 argparse to a very wide width and do the wrapping in post-processing.
107 :literal:`\\`\\`code\\`\\`` is replaced with just ``code`` in normal
108 output. :literal:`\\`\\`\\`code\\`\\`\\`` is replaced with ``'code'`` in
109 normal output and with :literal:`\\`\\`code\\`\\`` in ``rst`` output.
110 ``rst`` output is selected with environment variable
111 ``PYROCKO_RST_HELP=1``. The script ``maintenance/argparse_help_to_rst.py``
112 extracts the ``rst`` help and generates the ``rst`` files for the docs.
113 '''
115 def __init__(
116 self, prog=None, usage=None, description=None, epilog=None,
117 **kwargs):
119 kwargs['formatter_class'] = formatter_with_width(1000000)
121 description = unwrap(description)
122 epilog = unwrap(epilog)
124 argparse.ArgumentParser.__init__(
125 self, prog=prog, usage=usage, description=description,
126 epilog=epilog, **kwargs)
128 if hasattr(self, '_action_groups'):
129 for group in self._action_groups:
130 if group.title == 'positional arguments':
131 group.title = 'Positional arguments'
133 elif group.title == 'optional arguments':
134 group.title = 'Optional arguments'
136 elif group.title == 'options':
137 group.title = 'Options'
139 self.raw_help = False
141 def format_help(self, *args, **kwargs):
142 s = argparse.ArgumentParser.format_help(self, *args, **kwargs)
144 # replace usage with wrapped one from argparse because naive wrapping
145 # does not look good.
146 formatter_class = self.formatter_class
147 self.formatter_class = formatter_with_width(79)
148 usage = self.format_usage()
149 self.formatter_class = formatter_class
151 lines = []
152 for line in s.splitlines():
153 if line.startswith('usage:'):
154 lines.append(usage)
155 else:
156 lines.append(line)
158 s = '\n'.join(lines)
160 if os.environ.get('PYROCKO_RST_HELP', '0') == '0':
161 s = s.replace('```', "'")
162 s = s.replace('``', '')
163 s = wrap(s)
164 else:
165 s = s.replace('```', '``')
166 s = wrap_usage(s)
168 return s
171class SquirrelArgumentParser(PyrockoArgumentParser):
172 '''
173 Parser for CLI arguments with a some extras for Squirrel based apps.
175 :param command:
176 Implementation of the command.
177 :type command:
178 :py:class:`SquirrelCommand` or module providing the same interface
180 :param subcommands:
181 Implementations of subcommands.
182 :type subcommands:
183 :py:class:`list` of :py:class:`SquirrelCommand` or modules providing
184 the same interface
186 :param \\*args:
187 Handed through to base class's init.
189 :param \\*\\*kwargs:
190 Handed through to base class's init.
191 '''
193 def __init__(self, *args, command=None, subcommands=[], **kwargs):
195 self._command = command
196 self._subcommands = subcommands
197 self._have_selection_arguments = False
198 self._have_query_arguments = False
200 kwargs['add_help'] = False
201 PyrockoArgumentParser.__init__(self, *args, **kwargs)
202 add_standard_arguments(self)
203 self._command = None
204 self._subcommands = []
205 if command:
206 self.set_command(command)
208 if subcommands:
209 self.set_subcommands(subcommands)
211 def set_command(self, command):
212 command.setup(self)
213 self.set_defaults(target=command.run)
215 def set_subcommands(self, subcommands):
216 subparsers = self.add_subparsers(
217 metavar='SUBCOMMAND',
218 title='Subcommands')
220 for mod in subcommands:
221 subparser = mod.make_subparser(subparsers)
222 if subparser is None:
223 raise Exception(
224 'make_subparser(subparsers) must return the created '
225 'parser.')
227 mod.setup(subparser)
228 subparser.set_defaults(target=mod.run, subparser=subparser)
230 def parse_args(self, args=None, namespace=None):
231 '''
232 Parse arguments given on command line.
234 Extends the functionality of
235 :py:meth:`argparse.ArgumentParser.parse_args` to process and handle the
236 standard options ``--loglevel``, ``--progress`` and ``--help``.
237 '''
239 args = PyrockoArgumentParser.parse_args(
240 self, args=args, namespace=namespace)
242 eff_parser = args.__dict__.get('subparser', self)
244 process_standard_arguments(eff_parser, args)
246 if eff_parser._have_selection_arguments:
247 def make_squirrel():
248 return squirrel_from_selection_arguments(args)
250 args.make_squirrel = make_squirrel
252 if eff_parser._have_query_arguments:
253 try:
254 args.squirrel_query = squirrel_query_from_arguments(args)
255 except (error.SquirrelError, error.ToolError) as e:
256 logger.fatal(str(e))
257 sys.exit(1)
259 return args
261 def dispatch(self, args):
262 '''
263 Dispatch execution to selected command/subcommand.
265 :param args:
266 Parsed arguments obtained from :py:meth:`parse_args`.
268 :returns:
269 ``True`` if dispatching was successful, ``False`` otherwise.
271 If an exception of type
272 :py:exc:`~pyrocko.squirrel.error.SquirrelError` or
273 :py:exc:`~pyrocko.squirrel.error.ToolError` is caught, the error is
274 logged and the program is terminated with exit code 1.
275 '''
276 eff_parser = args.__dict__.get('subparser', self)
277 target = args.__dict__.get('target', None)
279 if target:
280 try:
281 target(eff_parser, args)
282 return True
284 except (error.SquirrelError, error.ToolError) as e:
285 logger.fatal(str(e))
286 sys.exit(1)
288 return False
290 def run(self, args=None):
291 '''
292 Parse arguments and dispatch to selected command/subcommand.
294 This simply calls :py:meth:`parse_args` and then :py:meth:`dispatch`
295 with the obtained ``args``. A usage message is printed if no command is
296 selected.
297 '''
298 args = self.parse_args(args)
299 if not self.dispatch(args):
300 self.print_help()
302 def add_squirrel_selection_arguments(self):
303 '''
304 Set up command line options commonly used to configure a
305 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
307 This will optional arguments ``--add``, ``--include``, ``--exclude``,
308 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``, and
309 ``--dataset``.
311 Call ``args.make_squirrel()`` on the arguments returned from
312 :py:meth:`parse_args` to finally instantiate and configure the
313 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
314 '''
315 add_squirrel_selection_arguments(self)
316 self._have_selection_arguments = True
318 def add_squirrel_query_arguments(self, without=[]):
319 '''
320 Set up command line options commonly used in squirrel queries.
322 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
323 ``--tmax``, and ``--time``.
325 Once finished with parsing, the query arguments are available as
326 ``args.squirrel_query`` on the arguments returned from
327 :py:meth:`parse_args`.
329 :param without:
330 Suppress adding given options.
331 :type without:
332 :py:class:`list` of :py:class:`str`, choices: ``'tmin'``,
333 ``'tmax'``, ``'codes'``, and ``'time'``.
334 '''
336 add_squirrel_query_arguments(self, without=without)
337 self._have_query_arguments = True
340def csvtype(choices):
341 def splitarg(arg):
342 values = arg.split(',')
343 for value in values:
344 if value not in choices:
345 raise argparse.ArgumentTypeError(
346 'Invalid choice: {!r} (choose from {})'
347 .format(value, ', '.join(map(repr, choices))))
348 return values
349 return splitarg
352def dq(x):
353 return '``%s``' % x
356def ldq(xs):
357 return ', '.join(dq(x) for x in xs)
360def add_standard_arguments(parser):
361 group = parser.add_argument_group('General options')
362 group.add_argument(
363 '--help', '-h',
364 action='help',
365 help='Show this help message and exit.')
367 loglevel_choices = ['critical', 'error', 'warning', 'info', 'debug']
368 loglevel_default = 'info'
370 group.add_argument(
371 '--loglevel',
372 choices=loglevel_choices,
373 default=loglevel_default,
374 metavar='LEVEL',
375 help='Set logger level. Choices: %s. Default: %s.' % (
376 ldq(loglevel_choices), dq(loglevel_default)))
378 progress_choices = ['terminal', 'log', 'off']
379 progress_default = 'terminal'
381 group.add_argument(
382 '--progress',
383 choices=progress_choices,
384 default=progress_default,
385 metavar='DEST',
386 help='Set how progress status is reported. Choices: %s. '
387 'Default: %s.' % (
388 ldq(progress_choices), dq(progress_default)))
391def process_standard_arguments(parser, args):
392 loglevel = args.__dict__.pop('loglevel')
393 util.setup_logging(parser.prog, loglevel)
395 pmode = args.__dict__.pop('progress')
396 progress.set_default_viewer(pmode)
399def add_squirrel_selection_arguments(parser):
400 '''
401 Set up command line options commonly used to configure a
402 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
404 This will optional arguments ``--add``, ``--include``, ``--exclude``,
405 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``,
406 and ``--dataset`` to a given argument parser.
408 Once finished with parsing, call
409 :py:func:`squirrel_from_selection_arguments` to finally instantiate and
410 configure the :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
412 :param parser:
413 The argument parser to be configured.
414 :type parser:
415 argparse.ArgumentParser
416 '''
417 from pyrocko import squirrel as sq
419 group = parser.add_argument_group('Data collection options')
421 group.add_argument(
422 '--add', '-a',
423 dest='paths',
424 metavar='PATH',
425 nargs='+',
426 help='Add files and directories with waveforms, metadata and events. '
427 'Content is indexed and added to the temporary (default) or '
428 'persistent (see ``--persistent``) data selection.')
430 group.add_argument(
431 '--include',
432 dest='include',
433 metavar='REGEX',
434 help='Only include files whose paths match the regular expression '
435 "``REGEX``. Examples: ``--include='\\.MSEED$'`` would only "
436 'match files ending with ```.MSEED```. '
437 "``--include='\\.BH[EN]\\.'`` would match paths containing "
438 "```.BHE.``` or ```.BHN.```. ``--include='/2011/'`` would "
439 'match paths with a subdirectory ```2011``` in their path '
440 'hierarchy.')
442 group.add_argument(
443 '--exclude',
444 dest='exclude',
445 metavar='REGEX',
446 help='Only include files whose paths do not match the regular '
447 "expression ``REGEX``. Examples: ``--exclude='/\\.DS_Store/'`` "
448 'would exclude anything inside any ```.DS_Store``` subdirectory.')
450 group.add_argument(
451 '--optimistic', '-o',
452 action='store_false',
453 dest='check',
454 default=True,
455 help='Disable checking file modification times for faster startup.')
457 group.add_argument(
458 '--format', '-f',
459 dest='format',
460 metavar='FORMAT',
461 default='detect',
462 choices=sq.supported_formats(),
463 help='Assume input files are of given ``FORMAT``. Choices: %s. '
464 'Default: %s.' % (
465 ldq(sq.supported_formats()),
466 dq('detect')))
468 group.add_argument(
469 '--add-only',
470 type=csvtype(sq.supported_content_kinds()),
471 dest='kinds_add',
472 metavar='KINDS',
473 help='Restrict meta-data scanning to given content kinds. '
474 '``KINDS`` is a comma-separated list of content kinds. '
475 'Choices: %s. By default, all content kinds are indexed.'
476 % ldq(sq.supported_content_kinds()))
478 group.add_argument(
479 '--persistent', '-p',
480 dest='persistent',
481 metavar='NAME',
482 help='Create/use persistent selection with given ``NAME``. Persistent '
483 'selections can be used to speed up startup of Squirrel-based '
484 'applications.')
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
517 squirrel = base.Squirrel(persistent=args.persistent)
519 if args.paths:
520 squirrel.add(
521 args.paths,
522 check=args.check,
523 format=args.format,
524 kinds=args.kinds_add or None,
525 include=args.include,
526 exclude=args.exclude)
528 for dataset_path in args.datasets:
529 squirrel.add_dataset(dataset_path, check=args.check)
531 return squirrel
534def add_squirrel_query_arguments(parser, without=[]):
535 '''
536 Set up command line options commonly used in squirrel queries.
538 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
539 ``--tmax``, and ``--time``.
541 Once finished with parsing, call
542 :py:func:`squirrel_query_from_arguments` to get the parsed values.
544 :param parser:
545 The argument parser to be configured.
546 :type parser:
547 argparse.ArgumentParser
549 :param without:
550 Suppress adding given options.
551 :type without:
552 :py:class:`list` of :py:class:`str`
553 '''
555 from pyrocko import squirrel as sq
557 group = parser.add_argument_group('Data query options')
559 if 'kinds' not in without:
560 group.add_argument(
561 '--kinds',
562 type=csvtype(sq.supported_content_kinds()),
563 dest='kinds',
564 metavar='KINDS',
565 help='Content kinds to query. ``KINDS`` is a comma-separated list '
566 'of content kinds. Choices: %s. By default, all content '
567 'kinds are queried.' % ldq(sq.supported_content_kinds()))
569 if 'codes' not in without:
570 group.add_argument(
571 '--codes',
572 dest='codes',
573 metavar='CODES',
574 help='Code patterns to query (``STA``, ``NET.STA``, '
575 '``NET.STA.LOC``, ``NET.STA.LOC.CHA``, or '
576 '``NET.STA.LOC.CHA.EXTRA``). The pattern may contain '
577 'wildcards ``*`` (zero or more arbitrary characters), ``?`` '
578 '(single arbitrary character), and ``[CHARS]`` (any '
579 'character out of ``CHARS``). Multiple patterns can be '
580 'given by separating them with commas.')
582 if 'tmin' not in without:
583 group.add_argument(
584 '--tmin',
585 dest='tmin',
586 metavar='TIME',
587 help='Begin of time interval to query. %s' % help_time_format)
589 if 'tmax' not in without:
590 group.add_argument(
591 '--tmax',
592 dest='tmax',
593 metavar='TIME',
594 help='End of time interval to query. %s' % help_time_format)
596 if 'time' not in without:
597 group.add_argument(
598 '--time',
599 dest='time',
600 metavar='TIME',
601 help='Time instant to query. %s' % help_time_format)
604def squirrel_query_from_arguments(args):
605 '''
606 Get common arguments to be used in squirrel queries from command line.
608 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
609 the necessary options.
611 :param args:
612 Parsed command line arguments, as returned by
613 :py:meth:`argparse.ArgumentParser.parse_args`.
615 :returns:
616 :py:class:`dict` with any parsed option values.
617 '''
619 from pyrocko import squirrel as sq
621 d = {}
623 if 'kinds' in args and args.kinds:
624 d['kind'] = args.kinds
625 if 'tmin' in args and args.tmin:
626 d['tmin'] = util.str_to_time_fillup(args.tmin)
627 if 'tmax' in args and args.tmax:
628 d['tmax'] = util.str_to_time_fillup(args.tmax)
629 if 'time' in args and args.time:
630 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
631 if 'codes' in args and args.codes:
632 d['codes'] = [
633 sq.to_codes_guess(s.strip()) for s in args.codes.split(',')]
635 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
636 raise error.SquirrelError(
637 'Options --tmin/--tmax and --time are mutually exclusive.')
638 return d
641class SquirrelCommand(object):
642 '''
643 Base class for Squirrel-based CLI programs and subcommands.
644 '''
646 def fail(self, message):
647 '''
648 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
650 :py:meth:`SquirrelArgumentParser.run` catches
651 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
652 terminates with an error exit state.
653 '''
654 raise error.ToolError(message)
656 def make_subparser(self, subparsers):
657 '''
658 To be implemented in subcommand. Create subcommand parser.
660 Must return a newly created parser obtained with
661 ``subparsers.add_parser(...)``, e.g.::
663 def make_subparser(self, subparsers):
664 return subparsers.add_parser(
665 'plot', help='Draw a nice plot.')
667 '''
668 return subparsers.add_parser(
669 self.__class__.__name__, help='Undocumented.')
671 def setup(self, parser):
672 '''
673 To be implemented in subcommand. Configure parser.
675 :param parser:
676 The argument parser to be configured.
677 :type parser:
678 argparse.ArgumentParser
680 Example::
682 def setup(self, parser):
683 parser.add_squirrel_selection_arguments()
684 parser.add_squirrel_query_arguments()
685 parser.add_argument(
686 '--fmin',
687 dest='fmin',
688 metavar='FLOAT',
689 type=float,
690 help='Corner of highpass [Hz].')
691 '''
692 pass
694 def run(self, parser, args):
695 '''
696 To be implemented in subcommand. Main routine of the command.
698 :param parser:
699 The argument parser to be configured.
700 :type parser:
701 argparse.ArgumentParser
703 :param args:
704 Parsed command line arguments, as returned by
705 :py:meth:`argparse.ArgumentParser.parse_args`.
707 Example::
709 def run(self, parser, args):
710 print('User has selected fmin = %g Hz' % args.fmin)
712 # args.make_squirrel() is available if
713 # parser.add_squirrel_selection_arguments() was called during
714 # setup().
716 sq = args.make_squirrel()
718 # args.squirrel_query is available if
719 # praser.add_squirrel_query_arguments() was called during
720 # setup().
722 stations = sq.get_stations(**args.squirrel_query)
723 '''
724 pass
727__all__ = [
728 'PyrockoArgumentParser',
729 'SquirrelArgumentParser',
730 'SquirrelCommand',
731 'add_squirrel_selection_arguments',
732 'squirrel_from_selection_arguments',
733 'add_squirrel_query_arguments',
734 'squirrel_query_from_arguments',
735]