1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function
8import sys
9import argparse
10import logging
12from pyrocko import util, progress
13from pyrocko.squirrel import error
16logger = logging.getLogger('psq.tool.common')
18help_time_format = 'Format: "YYYY-MM-DD HH:MM:SS.FFF", truncation allowed.'
21class PyrockoHelpFormatter(argparse.RawDescriptionHelpFormatter):
22 def __init__(self, *args, **kwargs):
23 kwargs['width'] = 79
24 argparse.RawDescriptionHelpFormatter.__init__(self, *args, **kwargs)
27class PyrockoArgumentParser(argparse.ArgumentParser):
29 def __init__(self, *args, **kwargs):
31 kwargs['formatter_class'] = PyrockoHelpFormatter
33 argparse.ArgumentParser.__init__(self, *args, **kwargs)
35 if hasattr(self, '_action_groups'):
36 for group in self._action_groups:
37 if group.title == 'positional arguments':
38 group.title = 'Positional arguments'
40 elif group.title == 'optional arguments':
41 group.title = 'Optional arguments'
44class SquirrelArgumentParser(PyrockoArgumentParser):
45 '''
46 Parser for CLI arguments with a some extras for Squirrel based apps.
48 :param command:
49 Implementation of the command.
50 :type command:
51 :py:class:`SquirrelCommand` (or module providing the same interface).
53 :param subcommands:
54 Implementations of subcommands.
55 :type subcommands:
56 :py:class:`list` of :py:class:`SquirrelCommand` (or modules providing
57 the same interface).
59 :param \\*args:
60 Handed through to base class's init.
62 :param \\*\\*kwargs:
63 Handed through to base class's init.
64 '''
66 def __init__(self, *args, command=None, subcommands=[], **kwargs):
68 self._command = command
69 self._subcommands = subcommands
70 self._have_selection_arguments = False
71 self._have_query_arguments = False
73 kwargs['add_help'] = False
74 PyrockoArgumentParser.__init__(self, *args, **kwargs)
75 add_standard_arguments(self)
76 self._command = None
77 self._subcommands = []
78 if command:
79 self.set_command(command)
81 if subcommands:
82 self.set_subcommands(subcommands)
84 def set_command(self, command):
85 command.setup(self)
86 self.set_defaults(target=command.run)
88 def set_subcommands(self, subcommands):
89 subparsers = self.add_subparsers(
90 metavar='SUBCOMMAND',
91 title='Subcommands')
93 for mod in subcommands:
94 subparser = mod.make_subparser(subparsers)
95 if subparser is None:
96 raise Exception(
97 'make_subparser(subparsers) must return the created '
98 'parser.')
100 mod.setup(subparser)
101 subparser.set_defaults(target=mod.run, subparser=subparser)
103 def parse_args(self, args=None, namespace=None):
104 '''
105 Parse arguments given on command line.
107 Extends the functionality of
108 :py:meth:`argparse.ArgumentParser.parse_args` to process and handle the
109 standard options ``--loglevel``, ``--progress`` and ``--help``.
110 '''
112 args = PyrockoArgumentParser.parse_args(
113 self, args=args, namespace=namespace)
115 eff_parser = args.__dict__.get('subparser', self)
117 process_standard_arguments(self, args)
119 if eff_parser._have_selection_arguments:
120 def make_squirrel():
121 return squirrel_from_selection_arguments(args)
123 args.make_squirrel = make_squirrel
125 if eff_parser._have_query_arguments:
126 try:
127 args.squirrel_query = squirrel_query_from_arguments(args)
128 except (error.SquirrelError, error.ToolError) as e:
129 logger.fatal(str(e))
130 sys.exit(1)
132 return args
134 def dispatch(self, args):
135 '''
136 Dispatch execution to selected command/subcommand.
138 :param args:
139 Parsed arguments obtained from :py:meth:`parse_args`.
141 :returns:
142 ``True`` if dispatching was successful, ``False`` othewise.
144 If an exception of type
145 :py:exc:`~pyrocko.squirrel.error.SquirrelError` or
146 :py:exc:`~pyrocko.squirrel.error.ToolError` is caught, the error is
147 logged and the program is terminated with exit code 1.
148 '''
149 eff_parser = args.__dict__.get('subparser', self)
150 target = args.__dict__.get('target', None)
152 if target:
153 try:
154 target(eff_parser, args)
155 return True
157 except (error.SquirrelError, error.ToolError) as e:
158 logger.fatal(str(e))
159 sys.exit(1)
161 return False
163 def run(self, args=None):
164 '''
165 Parse arguments and dispatch to selected command/subcommand.
167 This simply calls :py:meth:`parse_args` and then :py:meth:`dispatch`
168 with the obtained ``args``. A usage message is printed if no command is
169 selected.
170 '''
171 args = self.parse_args(args)
172 if not self.dispatch(args):
173 self.print_help()
175 def add_squirrel_selection_arguments(self):
176 '''
177 Set up command line options commonly used to configure a
178 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
180 This will optional arguments ``--add``, ``--include``, ``--exclude``,
181 ``--optimistic``, ``--format``, ``--kind``, ``--persistent``,
182 ``--update``, and ``--kind`` to a given argument parser.
184 Call ``args.make_squirrel()`` on the arguments returned from
185 :py:meth:`parse_args` to finally instantiate and configure the
186 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
187 '''
188 add_squirrel_selection_arguments(self)
189 self._have_selection_arguments = True
191 def add_squirrel_query_arguments(self, without=[]):
192 '''
193 Set up command line options commonly used in squirrel queries.
195 This will add options ``--codes``, ``--tmin``, ``--tmax``, and
196 ``--time``.
198 Once finished with parsing, the query arguments are available as
199 ``args.squirrel_query`` on the arguments returned from
200 :py:meth:`prase_args`.
202 :param without:
203 Suppress adding given options.
204 :type without:
205 :py:class:`list` of :py:class:`str`, choices: ``'tmin'``,
206 ``'tmax'``, ``'codes'``, and ``'time'``.
207 '''
209 add_squirrel_query_arguments(self, without=without)
210 self._have_query_arguments = True
213def csvtype(choices):
214 def splitarg(arg):
215 values = arg.split(',')
216 for value in values:
217 if value not in choices:
218 raise argparse.ArgumentTypeError(
219 'Invalid choice: {!r} (choose from {})'
220 .format(value, ', '.join(map(repr, choices))))
221 return values
222 return splitarg
225def add_standard_arguments(parser):
226 group = parser.add_argument_group('General options')
227 group.add_argument(
228 '--help', '-h',
229 action='help',
230 help='Show this help message and exit.')
232 group.add_argument(
233 '--loglevel',
234 choices=['critical', 'error', 'warning', 'info', 'debug'],
235 default='info',
236 metavar='LEVEL',
237 help='Set logger level. Choices: %(choices)s. Default: %(default)s.')
239 group.add_argument(
240 '--progress',
241 choices=['terminal', 'log', 'off'],
242 default='terminal',
243 metavar='DEST',
244 help='Set how progress status is reported. Choices: %(choices)s. '
245 'Default: %(default)s.')
248def process_standard_arguments(parser, args):
249 loglevel = args.__dict__.pop('loglevel')
250 util.setup_logging(parser.prog, loglevel)
252 pmode = args.__dict__.pop('progress')
253 progress.set_default_viewer(pmode)
256def add_squirrel_selection_arguments(parser):
257 '''
258 Set up command line options commonly used to configure a
259 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
261 This will optional arguments ``--add``, ``--include``, ``--exclude``,
262 ``--optimistic``, ``--format``, ``--kind``, ``--persistent``, ``--update``,
263 and ``--kind`` to a given argument parser.
265 Once finished with parsing, call
266 :py:func:`squirrel_from_selection_arguments` to finally instantiate and
267 configure the :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
269 :param parser:
270 The argument parser to be configured.
271 :type parser:
272 argparse.ArgumentParser
273 '''
274 from pyrocko import squirrel as sq
276 group = parser.add_argument_group('Data collection options')
278 group.add_argument(
279 '--add', '-a',
280 dest='paths',
281 metavar='PATH',
282 nargs='+',
283 help='Add files and directories with waveforms, metadata and events. '
284 'Content is indexed and added to the temporary (default) or '
285 'persistent (see --persistent) data selection.')
287 group.add_argument(
288 '--include',
289 dest='include',
290 metavar='REGEX',
291 help='Only include files whose paths match the regular expression '
292 'REGEX. Examples: --include=\'\\.MSEED$\' would only match files '
293 'ending with ".MSEED". --include=\'\\.BH[EN]\\.\' would match '
294 'paths containing ".BHE." or ".BHN.". --include=\'/2011/\' would '
295 'match paths with a subdirectory "2011" in their path hierarchy.')
297 group.add_argument(
298 '--exclude',
299 dest='exclude',
300 metavar='REGEX',
301 help='Only include files whose paths do not match the regular '
302 'expression REGEX. Examples: --exclude=\'/\\.DS_Store/\' would '
303 'exclude anything inside any ".DS_Store" subdirectory.')
305 group.add_argument(
306 '--optimistic', '-o',
307 action='store_false',
308 dest='check',
309 default=True,
310 help='Disable checking file modification times for faster startup.')
312 group.add_argument(
313 '--format', '-f',
314 dest='format',
315 metavar='FORMAT',
316 default='detect',
317 choices=sq.supported_formats(),
318 help='Assume input files are of given FORMAT. Choices: %(choices)s. '
319 'Default: %(default)s.')
321 group.add_argument(
322 '--add-only',
323 type=csvtype(sq.supported_content_kinds()),
324 dest='kinds_add',
325 metavar='KINDS',
326 help='Restrict meta-data scanning to given content kinds. '
327 'KINDS is a comma-separated list of content kinds, choices: %s. '
328 'By default, all content kinds are indexed.'
329 % ', '.join(sq.supported_content_kinds()))
331 group.add_argument(
332 '--persistent', '-p',
333 dest='persistent',
334 metavar='NAME',
335 help='Create/use persistent selection with given NAME. Persistent '
336 'selections can be used to speed up startup of Squirrel-based '
337 'applications.')
339 group.add_argument(
340 '--update', '-u',
341 dest='update',
342 action='store_true',
343 default=False,
344 help='Allow adding paths and datasets to existing persistent '
345 'selection.')
347 group.add_argument(
348 '--dataset', '-d',
349 dest='datasets',
350 default=[],
351 action='append',
352 metavar='FILE',
353 help='Add files, directories and remote sources from dataset '
354 'description file. This option can be repeated to add multiple '
355 'datasets. Run `squirrel template` to obtain examples of dataset '
356 'description files.')
359def squirrel_from_selection_arguments(args):
360 '''
361 Create a :py:class:`~pyrocko.squirrel.base.Squirrel` instance from command
362 line arguments.
364 Use :py:func:`add_squirrel_selection_arguments` to configure the parser
365 with the necessary options.
367 :param args:
368 Parsed command line arguments, as returned by
369 :py:meth:`argparse.ArgumentParser.parse_args`.
371 :returns:
372 :py:class:`pyrocko.squirrel.base.Squirrel` instance with paths,
373 datasets and remote sources added.
375 '''
376 from pyrocko.squirrel import base, dataset
378 datasets = [
379 dataset.read_dataset(dataset_path) for dataset_path in args.datasets]
381 persistents = [ds.persistent or '' for ds in datasets if ds.persistent]
382 if args.persistent:
383 persistent = args.persistent
384 elif persistents:
385 persistent = persistents[0]
386 if not all(p == persistents for p in persistents[1:]):
387 raise error.SquirrelError(
388 'Given datasets specify different `persistent` settings.')
390 if persistent:
391 logger.info(
392 'Persistent selection requested by dataset: %s' % persistent)
393 else:
394 persistent = None
396 else:
397 persistent = None
399 squirrel = base.Squirrel(persistent=persistent)
401 if persistent and not squirrel.is_new():
402 if not args.update:
403 logger.info(
404 'Using existing persistent selection: %s' % persistent)
405 if args.paths or datasets:
406 logger.info(
407 'Avoiding dataset rescan. Use --update/-u to '
408 'rescan or add items to existing persistent selection.')
410 return squirrel
412 else:
413 logger.info(
414 'Updating existing persistent selection: %s' % persistent)
416 if args.paths:
417 squirrel.add(
418 args.paths,
419 check=args.check,
420 format=args.format,
421 kinds=args.kinds_add or None,
422 include=args.include,
423 exclude=args.exclude)
425 for ds in datasets:
426 squirrel.add_dataset(ds, check=args.check)
428 return squirrel
431def add_squirrel_query_arguments(parser, without=[]):
432 '''
433 Set up command line options commonly used in squirrel queries.
435 This will add options ``--codes``, ``--tmin``, ``--tmax``, and ``--time``.
437 Once finished with parsing, call
438 :py:func:`squirrel_query_from_arguments` to get the parsed values.
440 :param parser:
441 The argument parser to be configured.
442 :type parser:
443 argparse.ArgumentParser
445 :param without:
446 Suppress adding given options.
447 :type without:
448 :py:class:`list` of :py:class:`str`
449 '''
451 from pyrocko import squirrel as sq
453 group = parser.add_argument_group('Data query options')
455 if 'kinds' not in without:
456 group.add_argument(
457 '--kinds',
458 type=csvtype(sq.supported_content_kinds()),
459 dest='kinds',
460 metavar='KINDS',
461 help='Content kinds to query. KINDS is a comma-separated list of '
462 'content kinds, choices: %s. By default, all content kinds '
463 'are queried.' % ', '.join(sq.supported_content_kinds()))
465 if 'codes' not in without:
466 group.add_argument(
467 '--codes',
468 dest='codes',
469 metavar='CODES',
470 help='Code pattern to query (STA, NET.STA, NET.STA.LOC, '
471 'NET.STA.LOC.CHA, or NET.STA.LOC.CHA.EXTRA).')
473 if 'tmin' not in without:
474 group.add_argument(
475 '--tmin',
476 dest='tmin',
477 metavar='TIME',
478 help='Begin of time interval to query. %s' % help_time_format)
480 if 'tmax' not in without:
481 group.add_argument(
482 '--tmax',
483 dest='tmax',
484 metavar='TIME',
485 help='End of time interval to query. %s' % help_time_format)
487 if 'time' not in without:
488 group.add_argument(
489 '--time',
490 dest='time',
491 metavar='TIME',
492 help='Time instant to query. %s' % help_time_format)
495def squirrel_query_from_arguments(args):
496 '''
497 Get common arguments to be used in squirrel queries from command line.
499 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
500 the necessary options.
502 :param args:
503 Parsed command line arguments, as returned by
504 :py:meth:`argparse.ArgumentParser.parse_args`.
506 :returns:
507 :py:class:`dict` with any parsed option values.
508 '''
510 from pyrocko import squirrel as sq
512 d = {}
514 if 'kinds' in args and args.kinds:
515 d['kind'] = args.kinds
516 if 'tmin' in args and args.tmin:
517 d['tmin'] = util.str_to_time_fillup(args.tmin)
518 if 'tmax' in args and args.tmax:
519 d['tmax'] = util.str_to_time_fillup(args.tmax)
520 if 'time' in args and args.time:
521 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
522 if 'codes' in args and args.codes:
523 d['codes'] = sq.to_codes_guess(args.codes)
525 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
526 raise error.SquirrelError(
527 'Options --tmin/--tmax and --time are mutually exclusive.')
528 return d
531class SquirrelCommand(object):
532 '''
533 Base class for Squirrel-based CLI programs and subcommands.
534 '''
536 def fail(self, message):
537 '''
538 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
540 :py:func:`~pyrocko.squirrel.tool.from_command` catches
541 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
542 terminates with an error exit state.
543 '''
544 raise error.ToolError(message)
546 def make_subparser(self, subparsers):
547 '''
548 To be implemented in subcommand. Create subcommand parser.
550 Must return a newly created parser obtained with
551 :py:meth:`add_parser`, e.g.::
553 def make_subparser(self, subparsers):
554 return subparsers.add_parser(
555 'plot', help='Draw a nice plot.')
557 '''
558 return subparsers.add_parser(
559 self.__class__.__name__, help='Undocumented.')
561 def setup(self, parser):
562 '''
563 To be implemented in subcommand. Configure parser.
565 :param parser:
566 The argument parser to be configured.
567 :type parser:
568 argparse.ArgumentParser
570 Example::
572 def setup(self, parser):
573 parser.add_squirrel_selection_arguments()
574 parser.add_squirrel_query_arguments()
575 parser.add_argument(
576 '--fmin',
577 dest='fmin',
578 metavar='FLOAT',
579 type=float,
580 help='Corner of highpass [Hz].')
581 '''
582 pass
584 def run(self, parser, args):
585 '''
586 To be implemented in subcommand. Main routine of the command.
588 :param parser:
589 The argument parser to be configured.
590 :type parser:
591 argparse.ArgumentParser
593 :param args:
594 Parsed command line arguments, as returned by
595 :py:meth:`argparse.ArgumentParser.parse_args`.
597 Example::
599 def run(self, parser, args):
600 print('User has selected fmin = %g Hz' % args.fmin)
602 # args.make_squirrel() is available if
603 # parser.add_squirrel_selection_arguments() was called during
604 # setup().
606 sq = args.make_squirrel()
608 # args.squirrel_query is available if
609 # praser.add_squirrel_query_arguments() was called during
610 # setup().
612 stations = sq.get_stations(**args.squirrel_query)
613 '''
614 pass
617__all__ = [
618 'SquirrelArgumentParser',
619 'SquirrelCommand',
620 'add_squirrel_selection_arguments',
621 'squirrel_from_selection_arguments',
622 'add_squirrel_query_arguments',
623 'squirrel_query_from_arguments',
624]