Coverage for /usr/local/lib/python3.13/dist-packages/pyrocko/squirrel/tool/common.py: 79%
243 statements
« prev ^ index » next coverage.py v7.6.0, created at 2025-12-04 10:41 +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(check_have_data=True):
262 return squirrel_from_selection_arguments(
263 args, check_have_data=check_have_data)
265 args.make_squirrel = make_squirrel
267 if eff_parser._have_query_arguments:
268 try:
269 args.squirrel_query = squirrel_query_from_arguments(args)
270 except (error.SquirrelError, error.ToolError) as e:
271 logger.fatal(str(e))
272 sys.exit(1)
274 return args
276 def dispatch(self, args):
277 '''
278 Dispatch execution to selected command/subcommand.
280 :param args:
281 Parsed arguments obtained from :py:meth:`parse_args`.
283 :returns:
284 ``True`` if dispatching was successful, ``False`` otherwise.
286 If an exception of type
287 :py:exc:`~pyrocko.squirrel.error.SquirrelError` or
288 :py:exc:`~pyrocko.squirrel.error.ToolError` is caught, the error is
289 logged and the program is terminated with exit code 1.
290 '''
291 eff_parser = args.__dict__.get('subparser', self)
292 target = args.__dict__.get('target', None)
294 if target:
295 try:
296 target(eff_parser, args)
297 return True
299 except (error.SquirrelError, error.ToolError) as e:
300 logger.fatal(str(e))
301 sys.exit(1)
303 return False
305 def run(self, args=None):
306 '''
307 Parse arguments and dispatch to selected command/subcommand.
309 This simply calls :py:meth:`parse_args` and then :py:meth:`dispatch`
310 with the obtained ``args``. A usage message is printed if no command is
311 selected.
312 '''
313 args = self.parse_args(args)
314 if not self.dispatch(args):
315 self.print_help()
317 def add_squirrel_selection_arguments(self):
318 '''
319 Set up command line options commonly used to configure a
320 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
322 This will optional arguments ``--add``, ``--include``, ``--exclude``,
323 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``, and
324 ``--dataset``.
326 Call ``args.make_squirrel()`` on the arguments returned from
327 :py:meth:`parse_args` to finally instantiate and configure the
328 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
329 '''
330 add_squirrel_selection_arguments(self)
331 self._have_selection_arguments = True
333 def add_squirrel_query_arguments(self, without=[]):
334 '''
335 Set up command line options commonly used in squirrel queries.
337 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
338 ``--tmax``, and ``--time``.
340 Once finished with parsing, the query arguments are available as
341 ``args.squirrel_query`` on the arguments returned from
342 :py:meth:`parse_args`.
344 :param without:
345 Suppress adding given options.
346 :type without:
347 :py:class:`list` of :py:class:`str`, choices: ``'tmin'``,
348 ``'tmax'``, ``'codes'``, and ``'time'``.
349 '''
351 add_squirrel_query_arguments(self, without=without)
352 self._have_query_arguments = True
355def csvtype(choices):
356 def splitarg(arg):
357 values = arg.split(',')
358 for value in values:
359 if value not in choices:
360 raise argparse.ArgumentTypeError(
361 'Invalid choice: {!r} (choose from {})'
362 .format(value, ', '.join(map(repr, choices))))
363 return values
364 return splitarg
367def dq(x):
368 return '``%s``' % x
371def ldq(xs):
372 return ', '.join(dq(x) for x in xs)
375def add_standard_arguments(parser):
376 group = parser.add_argument_group('General options')
377 group.add_argument(
378 '--help', '-h',
379 action='help',
380 help='Show this help message and exit.')
382 loglevel_choices = ['critical', 'error', 'warning', 'info', 'debug']
383 loglevel_default = 'info'
385 group.add_argument(
386 '--loglevel',
387 choices=loglevel_choices,
388 default=loglevel_default,
389 metavar='LEVEL',
390 help='Set logger level. Choices: %s. Default: %s.' % (
391 ldq(loglevel_choices), dq(loglevel_default)))
393 progress_choices = ['terminal', 'log', 'off']
394 progress_default = 'terminal'
396 group.add_argument(
397 '--progress',
398 choices=progress_choices,
399 default=progress_default,
400 metavar='DEST',
401 help='Set how progress status is reported. Choices: %s. '
402 'Default: %s.' % (
403 ldq(progress_choices), dq(progress_default)))
405 group.add_argument(
406 '--threads-loading',
407 dest='n_threads',
408 type=int,
409 default=1,
410 metavar='INT',
411 help='Set number of threads used in file content loading/decoding. '
412 'Default: 1')
415def process_standard_arguments(parser, args):
416 loglevel = args.__dict__.pop('loglevel')
417 util.setup_logging(parser.prog, loglevel)
419 pmode = args.__dict__.pop('progress')
420 progress.set_default_viewer(pmode)
423def add_squirrel_selection_arguments(parser):
424 '''
425 Set up command line options commonly used to configure a
426 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
428 This will optional arguments ``--add``, ``--include``, ``--exclude``,
429 ``--optimistic``, ``--format``, ``--add-only``, ``--persistent``,
430 and ``--dataset`` to a given argument parser.
432 Once finished with parsing, call
433 :py:func:`squirrel_from_selection_arguments` to finally instantiate and
434 configure the :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
436 :param parser:
437 The argument parser to be configured.
438 :type parser:
439 argparse.ArgumentParser
440 '''
441 from pyrocko import squirrel as sq
443 group = parser.add_argument_group('Data collection options')
445 group.add_argument(
446 '--add', '-a',
447 dest='paths',
448 metavar='PATH',
449 nargs='+',
450 help='Add files and directories with waveforms, metadata and events. '
451 'Content is indexed and added to the temporary (default) or '
452 'persistent (see ``--persistent``) data selection.')
454 group.add_argument(
455 '--include',
456 dest='include',
457 metavar='REGEX',
458 help='Only include files whose paths match the regular expression '
459 "``REGEX``. Examples: ``--include='\\.MSEED$'`` would only "
460 'match files ending with ```.MSEED```. '
461 "``--include='\\.BH[EN]\\.'`` would match paths containing "
462 "```.BHE.``` or ```.BHN.```. ``--include='/2011/'`` would "
463 'match paths with a subdirectory ```2011``` in their path '
464 'hierarchy.')
466 group.add_argument(
467 '--exclude',
468 dest='exclude',
469 metavar='REGEX',
470 help='Only include files whose paths do not match the regular '
471 "expression ``REGEX``. Examples: ``--exclude='/\\.DS_Store/'`` "
472 'would exclude anything inside any ```.DS_Store``` subdirectory.')
474 group.add_argument(
475 '--optimistic', '-o',
476 action='store_false',
477 dest='check',
478 default=True,
479 help='Disable checking file modification times for faster startup.')
481 group.add_argument(
482 '--format', '-f',
483 dest='format',
484 metavar='FORMAT',
485 default='detect',
486 choices=sq.supported_formats(),
487 help='Assume input files are of given ``FORMAT``. Choices: %s. '
488 'Default: %s.' % (
489 ldq(sq.supported_formats()),
490 dq('detect')))
492 group.add_argument(
493 '--add-only',
494 type=csvtype(sq.supported_content_kinds()),
495 dest='kinds_add',
496 metavar='KINDS',
497 help='Restrict meta-data scanning to given content kinds. '
498 '``KINDS`` is a comma-separated list of content kinds. '
499 'Choices: %s. By default, all content kinds are indexed.'
500 % ldq(sq.supported_content_kinds()))
502 group.add_argument(
503 '--persistent', '-p',
504 dest='persistent',
505 metavar='NAME',
506 help='Create/use persistent selection with given ``NAME``. Persistent '
507 'selections can be used to speed up startup of Squirrel-based '
508 'applications.')
510 group.add_argument(
511 '--samples-block',
512 dest='n_samples_block',
513 type=int,
514 metavar='N',
515 default=100000,
516 help='When downloading data, use blocks with approximately ``N`` '
517 'samples. Default: 100000')
519 group.add_argument(
520 '--dataset', '-d',
521 dest='datasets',
522 default=[],
523 action='append',
524 metavar='FILE',
525 help='Add files, directories and remote sources from dataset '
526 'description file. This option can be repeated to add multiple '
527 'datasets. Run ```squirrel template``` to obtain examples of '
528 'dataset description files.')
530 group.add_argument(
531 '--upgrade-storage',
532 dest='upgrade_storage',
533 default=False,
534 action='store_true',
535 help='Upgrade storage layout of cached data to latest version.')
538def squirrel_from_selection_arguments(args, check_have_data=True):
539 '''
540 Create a :py:class:`~pyrocko.squirrel.base.Squirrel` instance from command
541 line arguments.
543 Use :py:func:`add_squirrel_selection_arguments` to configure the parser
544 with the necessary options.
546 :param args:
547 Parsed command line arguments, as returned by
548 :py:meth:`argparse.ArgumentParser.parse_args`.
550 :returns:
551 :py:class:`pyrocko.squirrel.base.Squirrel` instance with paths,
552 datasets and remote sources added.
554 '''
555 from pyrocko.squirrel import base
557 if check_have_data and not (
558 args.persistent or args.datasets or args.paths):
560 raise error.ToolError(
561 'No data. Use --add, --dataset, and/or --persistent to set up '
562 'data sources.')
564 squirrel = base.Squirrel(
565 persistent=args.persistent,
566 n_samples_block=args.n_samples_block,
567 n_threads=getattr(args, 'n_threads', 1))
569 with progress.view():
570 if args.paths:
571 squirrel.add(
572 args.paths,
573 check=args.check,
574 format=args.format,
575 kinds=args.kinds_add or None,
576 include=args.include,
577 exclude=args.exclude)
579 with progress.task('add datasets', logger=logger) as task:
580 for dataset_path in task(args.datasets):
581 squirrel.add_dataset(
582 dataset_path,
583 check=args.check,
584 upgrade=args.upgrade_storage)
586 return squirrel
589def add_squirrel_query_arguments(parser, without=[]):
590 '''
591 Set up command line options commonly used in squirrel queries.
593 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
594 ``--tmax``, and ``--time``.
596 Once finished with parsing, call
597 :py:func:`squirrel_query_from_arguments` to get the parsed values.
599 :param parser:
600 The argument parser to be configured.
601 :type parser:
602 argparse.ArgumentParser
604 :param without:
605 Suppress adding given options.
606 :type without:
607 :py:class:`list` of :py:class:`str`
608 '''
610 from pyrocko import squirrel as sq
612 group = parser.add_argument_group('Data query options')
614 if 'kinds' not in without:
615 group.add_argument(
616 '--kinds',
617 type=csvtype(sq.supported_content_kinds()),
618 dest='kinds',
619 metavar='KINDS',
620 help='Content kinds to query. ``KINDS`` is a comma-separated list '
621 'of content kinds. Choices: %s. By default, all content '
622 'kinds are queried.' % ldq(sq.supported_content_kinds()))
624 if 'codes' not in without:
625 group.add_argument(
626 '--codes',
627 dest='codes',
628 metavar='CODES',
629 help='Code patterns to query (``STA``, ``NET.STA``, '
630 '``NET.STA.LOC``, ``NET.STA.LOC.CHA``, or '
631 '``NET.STA.LOC.CHA.EXTRA``). The pattern may contain '
632 'wildcards ``*`` (zero or more arbitrary characters), ``?`` '
633 '(single arbitrary character), and ``[CHARS]`` (any '
634 'character out of ``CHARS``). Multiple patterns can be '
635 'given by separating them with commas.')
637 group.add_argument(
638 '--codes-exclude',
639 dest='codes_exclude',
640 metavar='CODES',
641 help='Code patterns to exclude. See ``--codes``.')
643 if 'tmin' not in without:
644 group.add_argument(
645 '--tmin',
646 dest='tmin',
647 metavar='TIME',
648 help='Begin of time interval to query. %s' % help_time_format)
650 if 'tmax' not in without:
651 group.add_argument(
652 '--tmax',
653 dest='tmax',
654 metavar='TIME',
655 help='End of time interval to query. %s' % help_time_format)
657 if 'time' not in without:
658 group.add_argument(
659 '--time',
660 dest='time',
661 metavar='TIME',
662 help='Time instant to query. %s' % help_time_format)
665def ce(prefix, func, *args):
666 '''
667 Convert Exceptions in to ToolError.
668 '''
669 try:
670 return func(*args)
671 except Exception as e:
672 raise error.ToolError('%s: %s' % (prefix, str(e)))
675def squirrel_query_from_arguments(args):
676 '''
677 Get common arguments to be used in squirrel queries from command line.
679 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
680 the necessary options.
682 :param args:
683 Parsed command line arguments, as returned by
684 :py:meth:`argparse.ArgumentParser.parse_args`.
686 :returns:
687 :py:class:`dict` with any parsed option values.
688 '''
690 from pyrocko import squirrel as sq
692 d = {}
694 if 'kinds' in args and args.kinds:
695 d['kind'] = args.kinds
696 if 'tmin' in args and args.tmin:
697 d['tmin'] = ce(
698 '--tmin', util.str_to_time_fillup, args.tmin)
699 if 'tmax' in args and args.tmax:
700 d['tmax'] = ce(
701 '--tmax', util.str_to_time_fillup, args.tmax)
702 if 'time' in args and args.time:
703 d['tmin'] = d['tmax'] = ce(
704 '--time', util.str_to_time_fillup, args.time)
705 if 'codes' in args and args.codes:
706 d['codes'] = [
707 sq.to_codes_guess(s.strip()) for s in args.codes.split(',')]
708 if 'codes_exclude' in args and args.codes_exclude:
709 d['codes_exclude'] = [
710 sq.to_codes_guess(s.strip())
711 for s in args.codes_exclude.split(',')]
713 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
714 raise error.ToolError(
715 'Options --tmin/--tmax and --time are mutually exclusive.')
716 return d
719class SquirrelCommand(object):
720 '''
721 Base class for Squirrel-based CLI programs and subcommands.
722 '''
724 def fail(self, message):
725 '''
726 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
728 :py:meth:`SquirrelArgumentParser.run` catches
729 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
730 terminates with an error exit state.
731 '''
732 raise error.ToolError(message)
734 def make_subparser(self, subparsers):
735 '''
736 To be implemented in subcommand. Create subcommand parser.
738 Must return a newly created parser obtained with
739 ``subparsers.add_parser(...)``, e.g.::
741 def make_subparser(self, subparsers):
742 return subparsers.add_parser(
743 'plot', help='Draw a nice plot.')
745 '''
746 return subparsers.add_parser(
747 self.__class__.__name__, help='Undocumented.')
749 def setup(self, parser):
750 '''
751 To be implemented in subcommand. Configure parser.
753 :param parser:
754 The argument parser to be configured.
755 :type parser:
756 argparse.ArgumentParser
758 Example::
760 def setup(self, parser):
761 parser.add_squirrel_selection_arguments()
762 parser.add_squirrel_query_arguments()
763 parser.add_argument(
764 '--fmin',
765 dest='fmin',
766 metavar='FLOAT',
767 type=float,
768 help='Corner of highpass [Hz].')
769 '''
770 pass
772 def run(self, parser, args):
773 '''
774 To be implemented in subcommand. Main routine of the command.
776 :param parser:
777 The argument parser to be configured.
778 :type parser:
779 argparse.ArgumentParser
781 :param args:
782 Parsed command line arguments, as returned by
783 :py:meth:`argparse.ArgumentParser.parse_args`.
785 Example::
787 def run(self, parser, args):
788 print('User has selected fmin = %g Hz' % args.fmin)
790 # args.make_squirrel() is available if
791 # parser.add_squirrel_selection_arguments() was called during
792 # setup().
794 sq = args.make_squirrel()
796 # args.squirrel_query is available if
797 # praser.add_squirrel_query_arguments() was called during
798 # setup().
800 stations = sq.get_stations(**args.squirrel_query)
801 '''
802 pass
805__all__ = [
806 'PyrockoArgumentParser',
807 'SquirrelArgumentParser',
808 'SquirrelCommand',
809 'add_squirrel_selection_arguments',
810 'squirrel_from_selection_arguments',
811 'add_squirrel_query_arguments',
812 'squirrel_query_from_arguments',
813]