Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/squirrel/tool/common.py: 79%
233 statements
« prev ^ index » next coverage.py v6.5.0, created at 2024-07-04 08:07 +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['epilog'] = kwargs.get('epilog', '''
203----
205Manual: https://pyrocko.org/docs/current/apps/squirrel
207Tutorial: https://pyrocko.org/docs/current/apps/squirrel/tutorial.html
209Examples: https://pyrocko.org/docs/current/apps/squirrel/manual.html#examples
211🐿️
212''')
214 kwargs['add_help'] = False
215 PyrockoArgumentParser.__init__(self, *args, **kwargs)
216 add_standard_arguments(self)
217 self._command = None
218 self._subcommands = []
219 if command:
220 self.set_command(command)
222 if subcommands:
223 self.set_subcommands(subcommands)
225 def set_command(self, command):
226 command.setup(self)
227 self.set_defaults(target=command.run)
229 def set_subcommands(self, subcommands):
230 subparsers = self.add_subparsers(
231 metavar='SUBCOMMAND',
232 title='Subcommands')
234 for mod in subcommands:
235 subparser = mod.make_subparser(subparsers)
236 if subparser is None:
237 raise Exception(
238 'make_subparser(subparsers) must return the created '
239 'parser.')
241 mod.setup(subparser)
242 subparser.set_defaults(target=mod.run, subparser=subparser)
244 def parse_args(self, args=None, namespace=None):
245 '''
246 Parse arguments given on command line.
248 Extends the functionality of
249 :py:meth:`argparse.ArgumentParser.parse_args` to process and handle the
250 standard options ``--loglevel``, ``--progress`` and ``--help``.
251 '''
253 args = PyrockoArgumentParser.parse_args(
254 self, args=args, namespace=namespace)
256 eff_parser = args.__dict__.get('subparser', self)
258 process_standard_arguments(eff_parser, args)
260 if eff_parser._have_selection_arguments:
261 def make_squirrel():
262 return squirrel_from_selection_arguments(args)
264 args.make_squirrel = make_squirrel
266 if eff_parser._have_query_arguments:
267 try:
268 args.squirrel_query = squirrel_query_from_arguments(args)
269 except (error.SquirrelError, error.ToolError) as e:
270 logger.fatal(str(e))
271 sys.exit(1)
273 return args
275 def dispatch(self, args):
276 '''
277 Dispatch execution to selected command/subcommand.
279 :param args:
280 Parsed arguments obtained from :py:meth:`parse_args`.
282 :returns:
283 ``True`` if dispatching was successful, ``False`` otherwise.
285 If an exception of type
286 :py:exc:`~pyrocko.squirrel.error.SquirrelError` or
287 :py:exc:`~pyrocko.squirrel.error.ToolError` is caught, the error is
288 logged and the program is terminated with exit code 1.
289 '''
290 eff_parser = args.__dict__.get('subparser', self)
291 target = args.__dict__.get('target', None)
293 if target:
294 try:
295 target(eff_parser, args)
296 return True
298 except (error.SquirrelError, error.ToolError) as e:
299 logger.fatal(str(e))
300 sys.exit(1)
302 return False
304 def run(self, args=None):
305 '''
306 Parse arguments and dispatch to selected command/subcommand.
308 This simply calls :py:meth:`parse_args` and then :py:meth:`dispatch`
309 with the obtained ``args``. A usage message is printed if no command is
310 selected.
311 '''
312 args = self.parse_args(args)
313 if not self.dispatch(args):
314 self.print_help()
316 def add_squirrel_selection_arguments(self):
317 '''
318 Set up command line options commonly used to configure a
319 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
321 This will optional arguments ``--add``, ``--include``, ``--exclude``,
322 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``, and
323 ``--dataset``.
325 Call ``args.make_squirrel()`` on the arguments returned from
326 :py:meth:`parse_args` to finally instantiate and configure the
327 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
328 '''
329 add_squirrel_selection_arguments(self)
330 self._have_selection_arguments = True
332 def add_squirrel_query_arguments(self, without=[]):
333 '''
334 Set up command line options commonly used in squirrel queries.
336 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
337 ``--tmax``, and ``--time``.
339 Once finished with parsing, the query arguments are available as
340 ``args.squirrel_query`` on the arguments returned from
341 :py:meth:`parse_args`.
343 :param without:
344 Suppress adding given options.
345 :type without:
346 :py:class:`list` of :py:class:`str`, choices: ``'tmin'``,
347 ``'tmax'``, ``'codes'``, and ``'time'``.
348 '''
350 add_squirrel_query_arguments(self, without=without)
351 self._have_query_arguments = True
354def csvtype(choices):
355 def splitarg(arg):
356 values = arg.split(',')
357 for value in values:
358 if value not in choices:
359 raise argparse.ArgumentTypeError(
360 'Invalid choice: {!r} (choose from {})'
361 .format(value, ', '.join(map(repr, choices))))
362 return values
363 return splitarg
366def dq(x):
367 return '``%s``' % x
370def ldq(xs):
371 return ', '.join(dq(x) for x in xs)
374def add_standard_arguments(parser):
375 group = parser.add_argument_group('General options')
376 group.add_argument(
377 '--help', '-h',
378 action='help',
379 help='Show this help message and exit.')
381 loglevel_choices = ['critical', 'error', 'warning', 'info', 'debug']
382 loglevel_default = 'info'
384 group.add_argument(
385 '--loglevel',
386 choices=loglevel_choices,
387 default=loglevel_default,
388 metavar='LEVEL',
389 help='Set logger level. Choices: %s. Default: %s.' % (
390 ldq(loglevel_choices), dq(loglevel_default)))
392 progress_choices = ['terminal', 'log', 'off']
393 progress_default = 'terminal'
395 group.add_argument(
396 '--progress',
397 choices=progress_choices,
398 default=progress_default,
399 metavar='DEST',
400 help='Set how progress status is reported. Choices: %s. '
401 'Default: %s.' % (
402 ldq(progress_choices), dq(progress_default)))
405def process_standard_arguments(parser, args):
406 loglevel = args.__dict__.pop('loglevel')
407 util.setup_logging(parser.prog, loglevel)
409 pmode = args.__dict__.pop('progress')
410 progress.set_default_viewer(pmode)
413def add_squirrel_selection_arguments(parser):
414 '''
415 Set up command line options commonly used to configure a
416 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
418 This will optional arguments ``--add``, ``--include``, ``--exclude``,
419 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``,
420 and ``--dataset`` to a given argument parser.
422 Once finished with parsing, call
423 :py:func:`squirrel_from_selection_arguments` to finally instantiate and
424 configure the :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
426 :param parser:
427 The argument parser to be configured.
428 :type parser:
429 argparse.ArgumentParser
430 '''
431 from pyrocko import squirrel as sq
433 group = parser.add_argument_group('Data collection options')
435 group.add_argument(
436 '--add', '-a',
437 dest='paths',
438 metavar='PATH',
439 nargs='+',
440 help='Add files and directories with waveforms, metadata and events. '
441 'Content is indexed and added to the temporary (default) or '
442 'persistent (see ``--persistent``) data selection.')
444 group.add_argument(
445 '--include',
446 dest='include',
447 metavar='REGEX',
448 help='Only include files whose paths match the regular expression '
449 "``REGEX``. Examples: ``--include='\\.MSEED$'`` would only "
450 'match files ending with ```.MSEED```. '
451 "``--include='\\.BH[EN]\\.'`` would match paths containing "
452 "```.BHE.``` or ```.BHN.```. ``--include='/2011/'`` would "
453 'match paths with a subdirectory ```2011``` in their path '
454 'hierarchy.')
456 group.add_argument(
457 '--exclude',
458 dest='exclude',
459 metavar='REGEX',
460 help='Only include files whose paths do not match the regular '
461 "expression ``REGEX``. Examples: ``--exclude='/\\.DS_Store/'`` "
462 'would exclude anything inside any ```.DS_Store``` subdirectory.')
464 group.add_argument(
465 '--optimistic', '-o',
466 action='store_false',
467 dest='check',
468 default=True,
469 help='Disable checking file modification times for faster startup.')
471 group.add_argument(
472 '--format', '-f',
473 dest='format',
474 metavar='FORMAT',
475 default='detect',
476 choices=sq.supported_formats(),
477 help='Assume input files are of given ``FORMAT``. Choices: %s. '
478 'Default: %s.' % (
479 ldq(sq.supported_formats()),
480 dq('detect')))
482 group.add_argument(
483 '--add-only',
484 type=csvtype(sq.supported_content_kinds()),
485 dest='kinds_add',
486 metavar='KINDS',
487 help='Restrict meta-data scanning to given content kinds. '
488 '``KINDS`` is a comma-separated list of content kinds. '
489 'Choices: %s. By default, all content kinds are indexed.'
490 % ldq(sq.supported_content_kinds()))
492 group.add_argument(
493 '--persistent', '-p',
494 dest='persistent',
495 metavar='NAME',
496 help='Create/use persistent selection with given ``NAME``. Persistent '
497 'selections can be used to speed up startup of Squirrel-based '
498 'applications.')
500 group.add_argument(
501 '--dataset', '-d',
502 dest='datasets',
503 default=[],
504 action='append',
505 metavar='FILE',
506 help='Add files, directories and remote sources from dataset '
507 'description file. This option can be repeated to add multiple '
508 'datasets. Run ```squirrel template``` to obtain examples of '
509 'dataset description files.')
512def squirrel_from_selection_arguments(args):
513 '''
514 Create a :py:class:`~pyrocko.squirrel.base.Squirrel` instance from command
515 line arguments.
517 Use :py:func:`add_squirrel_selection_arguments` to configure the parser
518 with the necessary options.
520 :param args:
521 Parsed command line arguments, as returned by
522 :py:meth:`argparse.ArgumentParser.parse_args`.
524 :returns:
525 :py:class:`pyrocko.squirrel.base.Squirrel` instance with paths,
526 datasets and remote sources added.
528 '''
529 from pyrocko.squirrel import base
531 squirrel = base.Squirrel(persistent=args.persistent)
533 with progress.view():
534 if args.paths:
535 squirrel.add(
536 args.paths,
537 check=args.check,
538 format=args.format,
539 kinds=args.kinds_add or None,
540 include=args.include,
541 exclude=args.exclude)
543 with progress.task('add datasets', logger=logger) as task:
544 for dataset_path in task(args.datasets):
545 squirrel.add_dataset(
546 dataset_path, check=args.check)
548 return squirrel
551def add_squirrel_query_arguments(parser, without=[]):
552 '''
553 Set up command line options commonly used in squirrel queries.
555 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
556 ``--tmax``, and ``--time``.
558 Once finished with parsing, call
559 :py:func:`squirrel_query_from_arguments` to get the parsed values.
561 :param parser:
562 The argument parser to be configured.
563 :type parser:
564 argparse.ArgumentParser
566 :param without:
567 Suppress adding given options.
568 :type without:
569 :py:class:`list` of :py:class:`str`
570 '''
572 from pyrocko import squirrel as sq
574 group = parser.add_argument_group('Data query options')
576 if 'kinds' not in without:
577 group.add_argument(
578 '--kinds',
579 type=csvtype(sq.supported_content_kinds()),
580 dest='kinds',
581 metavar='KINDS',
582 help='Content kinds to query. ``KINDS`` is a comma-separated list '
583 'of content kinds. Choices: %s. By default, all content '
584 'kinds are queried.' % ldq(sq.supported_content_kinds()))
586 if 'codes' not in without:
587 group.add_argument(
588 '--codes',
589 dest='codes',
590 metavar='CODES',
591 help='Code patterns to query (``STA``, ``NET.STA``, '
592 '``NET.STA.LOC``, ``NET.STA.LOC.CHA``, or '
593 '``NET.STA.LOC.CHA.EXTRA``). The pattern may contain '
594 'wildcards ``*`` (zero or more arbitrary characters), ``?`` '
595 '(single arbitrary character), and ``[CHARS]`` (any '
596 'character out of ``CHARS``). Multiple patterns can be '
597 'given by separating them with commas.')
599 if 'tmin' not in without:
600 group.add_argument(
601 '--tmin',
602 dest='tmin',
603 metavar='TIME',
604 help='Begin of time interval to query. %s' % help_time_format)
606 if 'tmax' not in without:
607 group.add_argument(
608 '--tmax',
609 dest='tmax',
610 metavar='TIME',
611 help='End of time interval to query. %s' % help_time_format)
613 if 'time' not in without:
614 group.add_argument(
615 '--time',
616 dest='time',
617 metavar='TIME',
618 help='Time instant to query. %s' % help_time_format)
621def squirrel_query_from_arguments(args):
622 '''
623 Get common arguments to be used in squirrel queries from command line.
625 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
626 the necessary options.
628 :param args:
629 Parsed command line arguments, as returned by
630 :py:meth:`argparse.ArgumentParser.parse_args`.
632 :returns:
633 :py:class:`dict` with any parsed option values.
634 '''
636 from pyrocko import squirrel as sq
638 d = {}
640 if 'kinds' in args and args.kinds:
641 d['kind'] = args.kinds
642 if 'tmin' in args and args.tmin:
643 d['tmin'] = util.str_to_time_fillup(args.tmin)
644 if 'tmax' in args and args.tmax:
645 d['tmax'] = util.str_to_time_fillup(args.tmax)
646 if 'time' in args and args.time:
647 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
648 if 'codes' in args and args.codes:
649 d['codes'] = [
650 sq.to_codes_guess(s.strip()) for s in args.codes.split(',')]
652 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
653 raise error.SquirrelError(
654 'Options --tmin/--tmax and --time are mutually exclusive.')
655 return d
658class SquirrelCommand(object):
659 '''
660 Base class for Squirrel-based CLI programs and subcommands.
661 '''
663 def fail(self, message):
664 '''
665 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
667 :py:meth:`SquirrelArgumentParser.run` catches
668 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
669 terminates with an error exit state.
670 '''
671 raise error.ToolError(message)
673 def make_subparser(self, subparsers):
674 '''
675 To be implemented in subcommand. Create subcommand parser.
677 Must return a newly created parser obtained with
678 ``subparsers.add_parser(...)``, e.g.::
680 def make_subparser(self, subparsers):
681 return subparsers.add_parser(
682 'plot', help='Draw a nice plot.')
684 '''
685 return subparsers.add_parser(
686 self.__class__.__name__, help='Undocumented.')
688 def setup(self, parser):
689 '''
690 To be implemented in subcommand. Configure parser.
692 :param parser:
693 The argument parser to be configured.
694 :type parser:
695 argparse.ArgumentParser
697 Example::
699 def setup(self, parser):
700 parser.add_squirrel_selection_arguments()
701 parser.add_squirrel_query_arguments()
702 parser.add_argument(
703 '--fmin',
704 dest='fmin',
705 metavar='FLOAT',
706 type=float,
707 help='Corner of highpass [Hz].')
708 '''
709 pass
711 def run(self, parser, args):
712 '''
713 To be implemented in subcommand. Main routine of the command.
715 :param parser:
716 The argument parser to be configured.
717 :type parser:
718 argparse.ArgumentParser
720 :param args:
721 Parsed command line arguments, as returned by
722 :py:meth:`argparse.ArgumentParser.parse_args`.
724 Example::
726 def run(self, parser, args):
727 print('User has selected fmin = %g Hz' % args.fmin)
729 # args.make_squirrel() is available if
730 # parser.add_squirrel_selection_arguments() was called during
731 # setup().
733 sq = args.make_squirrel()
735 # args.squirrel_query is available if
736 # praser.add_squirrel_query_arguments() was called during
737 # setup().
739 stations = sq.get_stations(**args.squirrel_query)
740 '''
741 pass
744__all__ = [
745 'PyrockoArgumentParser',
746 'SquirrelArgumentParser',
747 'SquirrelCommand',
748 'add_squirrel_selection_arguments',
749 'squirrel_from_selection_arguments',
750 'add_squirrel_query_arguments',
751 'squirrel_query_from_arguments',
752]