Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/squirrel/tool/common.py: 79%
232 statements
« prev ^ index » next coverage.py v6.5.0, created at 2024-03-07 11:54 +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 with progress.view():
520 if args.paths:
521 squirrel.add(
522 args.paths,
523 check=args.check,
524 format=args.format,
525 kinds=args.kinds_add or None,
526 include=args.include,
527 exclude=args.exclude)
529 with progress.task('add datasets', logger=logger) as task:
530 for dataset_path in task(args.datasets):
531 squirrel.add_dataset(
532 dataset_path, check=args.check)
534 return squirrel
537def add_squirrel_query_arguments(parser, without=[]):
538 '''
539 Set up command line options commonly used in squirrel queries.
541 This will add optional arguments ``--kinds``, ``--codes``, ``--tmin``,
542 ``--tmax``, and ``--time``.
544 Once finished with parsing, call
545 :py:func:`squirrel_query_from_arguments` to get the parsed values.
547 :param parser:
548 The argument parser to be configured.
549 :type parser:
550 argparse.ArgumentParser
552 :param without:
553 Suppress adding given options.
554 :type without:
555 :py:class:`list` of :py:class:`str`
556 '''
558 from pyrocko import squirrel as sq
560 group = parser.add_argument_group('Data query options')
562 if 'kinds' not in without:
563 group.add_argument(
564 '--kinds',
565 type=csvtype(sq.supported_content_kinds()),
566 dest='kinds',
567 metavar='KINDS',
568 help='Content kinds to query. ``KINDS`` is a comma-separated list '
569 'of content kinds. Choices: %s. By default, all content '
570 'kinds are queried.' % ldq(sq.supported_content_kinds()))
572 if 'codes' not in without:
573 group.add_argument(
574 '--codes',
575 dest='codes',
576 metavar='CODES',
577 help='Code patterns to query (``STA``, ``NET.STA``, '
578 '``NET.STA.LOC``, ``NET.STA.LOC.CHA``, or '
579 '``NET.STA.LOC.CHA.EXTRA``). The pattern may contain '
580 'wildcards ``*`` (zero or more arbitrary characters), ``?`` '
581 '(single arbitrary character), and ``[CHARS]`` (any '
582 'character out of ``CHARS``). Multiple patterns can be '
583 'given by separating them with commas.')
585 if 'tmin' not in without:
586 group.add_argument(
587 '--tmin',
588 dest='tmin',
589 metavar='TIME',
590 help='Begin of time interval to query. %s' % help_time_format)
592 if 'tmax' not in without:
593 group.add_argument(
594 '--tmax',
595 dest='tmax',
596 metavar='TIME',
597 help='End of time interval to query. %s' % help_time_format)
599 if 'time' not in without:
600 group.add_argument(
601 '--time',
602 dest='time',
603 metavar='TIME',
604 help='Time instant to query. %s' % help_time_format)
607def squirrel_query_from_arguments(args):
608 '''
609 Get common arguments to be used in squirrel queries from command line.
611 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
612 the necessary options.
614 :param args:
615 Parsed command line arguments, as returned by
616 :py:meth:`argparse.ArgumentParser.parse_args`.
618 :returns:
619 :py:class:`dict` with any parsed option values.
620 '''
622 from pyrocko import squirrel as sq
624 d = {}
626 if 'kinds' in args and args.kinds:
627 d['kind'] = args.kinds
628 if 'tmin' in args and args.tmin:
629 d['tmin'] = util.str_to_time_fillup(args.tmin)
630 if 'tmax' in args and args.tmax:
631 d['tmax'] = util.str_to_time_fillup(args.tmax)
632 if 'time' in args and args.time:
633 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
634 if 'codes' in args and args.codes:
635 d['codes'] = [
636 sq.to_codes_guess(s.strip()) for s in args.codes.split(',')]
638 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
639 raise error.SquirrelError(
640 'Options --tmin/--tmax and --time are mutually exclusive.')
641 return d
644class SquirrelCommand(object):
645 '''
646 Base class for Squirrel-based CLI programs and subcommands.
647 '''
649 def fail(self, message):
650 '''
651 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
653 :py:meth:`SquirrelArgumentParser.run` catches
654 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
655 terminates with an error exit state.
656 '''
657 raise error.ToolError(message)
659 def make_subparser(self, subparsers):
660 '''
661 To be implemented in subcommand. Create subcommand parser.
663 Must return a newly created parser obtained with
664 ``subparsers.add_parser(...)``, e.g.::
666 def make_subparser(self, subparsers):
667 return subparsers.add_parser(
668 'plot', help='Draw a nice plot.')
670 '''
671 return subparsers.add_parser(
672 self.__class__.__name__, help='Undocumented.')
674 def setup(self, parser):
675 '''
676 To be implemented in subcommand. Configure parser.
678 :param parser:
679 The argument parser to be configured.
680 :type parser:
681 argparse.ArgumentParser
683 Example::
685 def setup(self, parser):
686 parser.add_squirrel_selection_arguments()
687 parser.add_squirrel_query_arguments()
688 parser.add_argument(
689 '--fmin',
690 dest='fmin',
691 metavar='FLOAT',
692 type=float,
693 help='Corner of highpass [Hz].')
694 '''
695 pass
697 def run(self, parser, args):
698 '''
699 To be implemented in subcommand. Main routine of the command.
701 :param parser:
702 The argument parser to be configured.
703 :type parser:
704 argparse.ArgumentParser
706 :param args:
707 Parsed command line arguments, as returned by
708 :py:meth:`argparse.ArgumentParser.parse_args`.
710 Example::
712 def run(self, parser, args):
713 print('User has selected fmin = %g Hz' % args.fmin)
715 # args.make_squirrel() is available if
716 # parser.add_squirrel_selection_arguments() was called during
717 # setup().
719 sq = args.make_squirrel()
721 # args.squirrel_query is available if
722 # praser.add_squirrel_query_arguments() was called during
723 # setup().
725 stations = sq.get_stations(**args.squirrel_query)
726 '''
727 pass
730__all__ = [
731 'PyrockoArgumentParser',
732 'SquirrelArgumentParser',
733 'SquirrelCommand',
734 'add_squirrel_selection_arguments',
735 'squirrel_from_selection_arguments',
736 'add_squirrel_query_arguments',
737 'squirrel_query_from_arguments',
738]