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 def __init__(self, *args, command=None, subcommands=[], **kwargs):
47 self._command = command
48 self._subcommands = subcommands
49 self._have_selection_arguments = False
50 self._have_query_arguments = False
52 kwargs['add_help'] = False
53 PyrockoArgumentParser.__init__(self, *args, **kwargs)
54 add_standard_arguments(self)
55 self._command = None
56 self._subcommands = []
57 if command:
58 self.set_command(command)
60 if subcommands:
61 self.set_subcommands(subcommands)
63 def set_command(self, command):
64 command.setup(self)
65 self.set_defaults(target=command.run)
67 def set_subcommands(self, subcommands):
68 subparsers = self.add_subparsers(
69 metavar='SUBCOMMAND',
70 title='Subcommands')
72 for mod in subcommands:
73 subparser = mod.make_subparser(subparsers)
74 if subparser is None:
75 raise Exception(
76 'make_subparser(subparsers) must return the created '
77 'parser.')
79 mod.setup(subparser)
80 subparser.set_defaults(target=mod.run, subparser=subparser)
82 def parse_args(self, args=None, namespace=None):
84 args = PyrockoArgumentParser.parse_args(
85 self, args=args, namespace=namespace)
87 eff_parser = args.__dict__.get('subparser', self)
89 process_standard_arguments(self, args)
91 if eff_parser._have_selection_arguments:
92 def make_squirrel():
93 return squirrel_from_selection_arguments(args)
95 args.make_squirrel = make_squirrel
97 if eff_parser._have_query_arguments:
98 try:
99 args.squirrel_query = squirrel_query_from_arguments(args)
100 except (error.SquirrelError, error.ToolError) as e:
101 logger.fatal(str(e))
102 sys.exit(1)
104 return args
106 def dispatch(self, args):
107 eff_parser = args.__dict__.get('subparser', self)
108 target = args.__dict__.get('target', None)
110 if target:
111 try:
112 target(eff_parser, args)
113 return True
115 except (error.SquirrelError, error.ToolError) as e:
116 logger.fatal(str(e))
117 sys.exit(1)
119 return False
121 def run(self, args=None):
122 args = self.parse_args(args)
123 if not self.dispatch(args):
124 self.print_help()
126 def add_squirrel_selection_arguments(self):
127 add_squirrel_selection_arguments(self)
128 self._have_selection_arguments = True
130 def add_squirrel_query_arguments(self, without=[]):
131 add_squirrel_query_arguments(self, without=without)
132 self._have_query_arguments = True
135def csvtype(choices):
136 def splitarg(arg):
137 values = arg.split(',')
138 for value in values:
139 if value not in choices:
140 raise argparse.ArgumentTypeError(
141 'Invalid choice: {!r} (choose from {})'
142 .format(value, ', '.join(map(repr, choices))))
143 return values
144 return splitarg
147def add_standard_arguments(parser):
148 group = parser.add_argument_group('General options')
149 group.add_argument(
150 '--help', '-h',
151 action='help',
152 help='Show this help message and exit.')
154 group.add_argument(
155 '--loglevel',
156 choices=['critical', 'error', 'warning', 'info', 'debug'],
157 default='info',
158 metavar='LEVEL',
159 help='Set logger level. Choices: %(choices)s. Default: %(default)s.')
161 group.add_argument(
162 '--progress',
163 choices=['terminal', 'log', 'off'],
164 default='terminal',
165 metavar='DEST',
166 help='Set how progress status is reported. Choices: %(choices)s. '
167 'Default: %(default)s.')
170def process_standard_arguments(parser, args):
171 loglevel = args.__dict__.pop('loglevel')
172 util.setup_logging(parser.prog, loglevel)
174 pmode = args.__dict__.pop('progress')
175 progress.set_default_viewer(pmode)
178def add_squirrel_selection_arguments(parser):
179 '''
180 Set up command line options commonly used to configure a
181 :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
183 This will optional arguments ``--add``, ``--include``, ``--exclude``,
184 ``--optimistic``, ``--format``, ``--kind``, ``--persistent``, ``--update``,
185 and ``--kind`` to a given argument parser.
187 Once finished with parsing, call
188 :py:func:`squirrel_from_selection_arguments` to finally instantiate and
189 configure the :py:class:`~pyrocko.squirrel.base.Squirrel` instance.
191 :param parser:
192 The argument parser to be configured.
193 :type parser:
194 argparse.ArgumentParser
195 '''
196 from pyrocko import squirrel as sq
198 group = parser.add_argument_group('Data collection options')
200 group.add_argument(
201 '--add', '-a',
202 dest='paths',
203 metavar='PATH',
204 nargs='+',
205 help='Add files and directories with waveforms, metadata and events. '
206 'Content is indexed and added to the temporary (default) or '
207 'persistent (see --persistent) data selection.')
209 group.add_argument(
210 '--include',
211 dest='include',
212 metavar='REGEX',
213 help='Only include files whose paths match the regular expression '
214 'REGEX. Examples: --include=\'\\.MSEED$\' would only match files '
215 'ending with ".MSEED". --include=\'\\.BH[EN]\\.\' would match '
216 'paths containing ".BHE." or ".BHN.". --include=\'/2011/\' would '
217 'match paths with a subdirectory "2011" in their path hierarchy.')
219 group.add_argument(
220 '--exclude',
221 dest='exclude',
222 metavar='REGEX',
223 help='Only include files whose paths do not match the regular '
224 'expression REGEX. Examples: --exclude=\'/\\.DS_Store/\' would '
225 'exclude anything inside any ".DS_Store" subdirectory.')
227 group.add_argument(
228 '--optimistic', '-o',
229 action='store_false',
230 dest='check',
231 default=True,
232 help='Disable checking file modification times for faster startup.')
234 group.add_argument(
235 '--format', '-f',
236 dest='format',
237 metavar='FORMAT',
238 default='detect',
239 choices=sq.supported_formats(),
240 help='Assume input files are of given FORMAT. Choices: %(choices)s. '
241 'Default: %(default)s.')
243 group.add_argument(
244 '--add-only',
245 type=csvtype(sq.supported_content_kinds()),
246 dest='kinds_add',
247 metavar='KINDS',
248 help='Restrict meta-data scanning to given content kinds. '
249 'KINDS is a comma-separated list of content kinds, choices: %s. '
250 'By default, all content kinds are indexed.'
251 % ', '.join(sq.supported_content_kinds()))
253 group.add_argument(
254 '--persistent', '-p',
255 dest='persistent',
256 metavar='NAME',
257 help='Create/use persistent selection with given NAME. Persistent '
258 'selections can be used to speed up startup of Squirrel-based '
259 'applications.')
261 group.add_argument(
262 '--update', '-u',
263 dest='update',
264 action='store_true',
265 default=False,
266 help='Allow adding paths and datasets to existing persistent '
267 'selection.')
269 group.add_argument(
270 '--dataset', '-d',
271 dest='datasets',
272 default=[],
273 action='append',
274 metavar='FILE',
275 help='Add files, directories and remote sources from dataset '
276 'description file. This option can be repeated to add multiple '
277 'datasets. Run `squirrel template` to obtain examples of dataset '
278 'description files.')
281def squirrel_from_selection_arguments(args):
282 '''
283 Create a :py:class:`~pyrocko.squirrel.base.Squirrel` instance from command
284 line arguments.
286 Use :py:func:`add_squirrel_selection_arguments` to configure the parser
287 with the necessary options.
289 :param args:
290 Parsed command line arguments, as returned by
291 :py:meth:`argparse.ArgumentParser.parse_args`.
293 :returns:
294 :py:class:`pyrocko.squirrel.base.Squirrel` instance with paths,
295 datasets and remote sources added.
297 '''
298 from pyrocko.squirrel import base, dataset
300 datasets = [
301 dataset.read_dataset(dataset_path) for dataset_path in args.datasets]
303 persistents = [ds.persistent or '' for ds in datasets if ds.persistent]
304 if args.persistent:
305 persistent = args.persistent
306 elif persistents:
307 persistent = persistents[0]
308 if not all(p == persistents for p in persistents[1:]):
309 raise error.SquirrelError(
310 'Given datasets specify different `persistent` settings.')
312 if persistent:
313 logger.info(
314 'Persistent selection requested by dataset: %s' % persistent)
315 else:
316 persistent = None
318 else:
319 persistent = None
321 squirrel = base.Squirrel(persistent=persistent)
323 if persistent and not squirrel.is_new():
324 if not args.update:
325 logger.info(
326 'Using existing persistent selection: %s' % persistent)
327 if args.paths or datasets:
328 logger.info(
329 'Avoiding dataset rescan. Use --update/-u to '
330 'rescan or add items to existing persistent selection.')
332 return squirrel
334 else:
335 logger.info(
336 'Updating existing persistent selection: %s' % persistent)
338 if args.paths:
339 squirrel.add(
340 args.paths,
341 check=args.check,
342 format=args.format,
343 kinds=args.kinds_add or None,
344 include=args.include,
345 exclude=args.exclude)
347 for ds in datasets:
348 squirrel.add_dataset(ds, check=args.check)
350 return squirrel
353def add_squirrel_query_arguments(parser, without=[]):
354 '''
355 Set up command line options commonly used in squirrel queries.
357 This will add options ``--codes``, ``--tmin``, ``--tmax``, and ``--time``.
359 Once finished with parsing, call
360 :py:func:`squirrel_query_from_arguments` to get the parsed values.
362 :param parser:
363 The argument parser to be configured.
364 :type parser:
365 argparse.ArgumentParser
367 :param without:
368 Suppress adding given options.
369 :type without:
370 :py:class:`list` of :py:class:`str`
371 '''
373 from pyrocko import squirrel as sq
375 group = parser.add_argument_group('Data query options')
377 if 'kinds' not in without:
378 group.add_argument(
379 '--kinds',
380 type=csvtype(sq.supported_content_kinds()),
381 dest='kinds',
382 metavar='KINDS',
383 help='Content kinds to query. KINDS is a comma-separated list of '
384 'content kinds, choices: %s. By default, all content kinds '
385 'are queried.' % ', '.join(sq.supported_content_kinds()))
387 if 'codes' not in without:
388 group.add_argument(
389 '--codes',
390 dest='codes',
391 metavar='CODES',
392 help='Code pattern to query (STA, NET.STA, NET.STA.LOC, '
393 'NET.STA.LOC.CHA, or NET.STA.LOC.CHA.EXTRA).')
395 if 'tmin' not in without:
396 group.add_argument(
397 '--tmin',
398 dest='tmin',
399 metavar='TIME',
400 help='Begin of time interval to query. %s' % help_time_format)
402 if 'tmax' not in without:
403 group.add_argument(
404 '--tmax',
405 dest='tmax',
406 metavar='TIME',
407 help='End of time interval to query. %s' % help_time_format)
409 if 'time' not in without:
410 group.add_argument(
411 '--time',
412 dest='time',
413 metavar='TIME',
414 help='Time instant to query. %s' % help_time_format)
417def squirrel_query_from_arguments(args):
418 '''
419 Get common arguments to be used in squirrel queries from command line.
421 Use :py:func:`add_squirrel_query_arguments` to configure the parser with
422 the necessary options.
424 :param args:
425 Parsed command line arguments, as returned by
426 :py:meth:`argparse.ArgumentParser.parse_args`.
428 :returns:
429 :py:class:`dict` with any parsed option values.
430 '''
432 from pyrocko import squirrel as sq
434 d = {}
436 if 'kinds' in args and args.kinds:
437 d['kind'] = args.kinds
438 if 'tmin' in args and args.tmin:
439 d['tmin'] = util.str_to_time_fillup(args.tmin)
440 if 'tmax' in args and args.tmax:
441 d['tmax'] = util.str_to_time_fillup(args.tmax)
442 if 'time' in args and args.time:
443 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
444 if 'codes' in args and args.codes:
445 d['codes'] = sq.to_codes_guess(args.codes)
447 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
448 raise error.SquirrelError(
449 'Options --tmin/--tmax and --time are mutually exclusive.')
450 return d
453class SquirrelCommand(object):
454 '''
455 Base class for Squirrel-based CLI programs and subcommands.
456 '''
458 def fail(self, message):
459 '''
460 Raises :py:exc:`~pyrocko.squirrel.error.ToolError`.
462 :py:func:`~pyrocko.squirrel.tool.from_command` catches
463 :py:exc:`~pyrocko.squirrel.error.ToolError`, logs the error message and
464 terminates with an error exit state.
465 '''
466 raise error.ToolError(message)
468 def make_subparser(self, subparsers):
469 '''
470 To be implemented in subcommand. Create subcommand parser.
472 Must return a newly created parser obtained with
473 :py:meth:`add_parser`, e.g.::
475 def make_subparser(self, subparsers):
476 return subparsers.add_parser(
477 'plot', help='Draw a nice plot.')
479 '''
480 return subparsers.add_parser(
481 self.__class__.__name__, help='Undocumented.')
483 def setup(self, parser):
484 '''
485 To be implemented in subcommand. Configure parser.
487 :param parser:
488 The argument parser to be configured.
489 :type parser:
490 argparse.ArgumentParser
492 Example::
494 def setup(self, parser):
495 parser.add_squirrel_selection_arguments()
496 parser.add_squirrel_query_arguments()
497 parser.add_argument(
498 '--fmin',
499 dest='fmin',
500 metavar='FLOAT',
501 type=float,
502 help='Corner of highpass [Hz].')
503 '''
504 pass
506 def run(self, parser, args):
507 '''
508 To be implemented in subcommand. Main routine of the command.
510 :param parser:
511 The argument parser to be configured.
512 :type parser:
513 argparse.ArgumentParser
515 :param args:
516 Parsed command line arguments, as returned by
517 :py:meth:`argparse.ArgumentParser.parse_args`.
519 Example::
521 def run(self, parser, args):
522 print('User has selected fmin = %g Hz' % args.fmin)
524 # args.make_squirrel() is available if
525 # parser.add_squirrel_selection_arguments() was called during
526 # setup().
528 sq = args.make_squirrel()
530 # args.squirrel_query is available if
531 # praser.add_squirrel_query_arguments() was called during
532 # setup().
534 stations = sq.get_stations(**args.squirrel_query)
535 '''
536 pass
539__all__ = [
540 'SquirrelArgumentParser',
541 'SquirrelCommand',
542 'add_squirrel_selection_arguments',
543 'squirrel_from_selection_arguments',
544 'add_squirrel_query_arguments',
545 'squirrel_query_from_arguments',
546]