1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function
8import argparse
9import logging
11from pyrocko import util, progress
12from pyrocko.squirrel import error
15logger = logging.getLogger('psq.tool.common')
18class PyrockoHelpFormatter(argparse.RawDescriptionHelpFormatter):
19 def __init__(self, *args, **kwargs):
20 kwargs['width'] = 79
21 argparse.RawDescriptionHelpFormatter.__init__(self, *args, **kwargs)
24class PyrockoArgumentParser(argparse.ArgumentParser):
26 def __init__(self, *args, **kwargs):
28 kwargs['formatter_class'] = PyrockoHelpFormatter
30 argparse.ArgumentParser.__init__(self, *args, **kwargs)
32 if hasattr(self, '_action_groups'):
33 for group in self._action_groups:
34 if group.title == 'positional arguments':
35 group.title = 'Positional arguments'
37 elif group.title == 'optional arguments':
38 group.title = 'Optional arguments'
41def csvtype(choices):
42 def splitarg(arg):
43 values = arg.split(',')
44 for value in values:
45 if value not in choices:
46 raise argparse.ArgumentTypeError(
47 'Invalid choice: {!r} (choose from {})'
48 .format(value, ', '.join(map(repr, choices))))
49 return values
50 return splitarg
53def add_parser(subparsers, *args, **kwargs):
54 kwargs['add_help'] = False
55 p = subparsers.add_parser(*args, **kwargs)
56 add_standard_arguments(p)
57 return p
60def add_standard_arguments(p):
61 p.add_argument(
62 '--help', '-h',
63 action='help',
64 help='Show this help message and exit.')
66 p.add_argument(
67 '--loglevel',
68 choices=['critical', 'error', 'warning', 'info', 'debug'],
69 default='info',
70 metavar='LEVEL',
71 help='Set logger level. Choices: %(choices)s. Default: %(default)s.')
73 p.add_argument(
74 '--progress',
75 choices=['terminal', 'log', 'off'],
76 default='terminal',
77 metavar='DEST',
78 help='Set how progress status is reported. Choices: %(choices)s. '
79 'Default: %(default)s.')
82def process_standard_arguments(parser, args):
83 loglevel = args.__dict__.pop('loglevel')
84 util.setup_logging(parser.prog, loglevel)
86 pmode = args.__dict__.pop('progress')
87 progress.set_default_viewer(pmode)
90def add_selection_arguments(p):
91 from pyrocko import squirrel as sq
93 p.add_argument(
94 'paths',
95 nargs='*',
96 help='Files and directories with waveforms, metadata and events.')
98 p.add_argument(
99 '--include',
100 dest='include',
101 metavar='REGEX',
102 help='Only include files whose paths match REGEX. Examples: '
103 '--include=\'\\.MSEED$\' would only match files ending with '
104 '".MSEED". --include=\'\\.BH[EN]\\.\' would match paths '
105 'containing ".BHE." or ".BHN.". --include=\'/2011/\' would match '
106 'paths with a subdirectory "2011" in their path hierarchy.')
108 p.add_argument(
109 '--exclude',
110 dest='exclude',
111 metavar='REGEX',
112 help='If given, files are only included if their paths do not match '
113 'the regular expression pattern REGEX.')
115 p.add_argument(
116 '--optimistic', '-o',
117 action='store_false',
118 dest='check',
119 default=True,
120 help='Disable checking file modification times.')
122 p.add_argument(
123 '--format', '-f',
124 dest='format',
125 metavar='FORMAT',
126 default='detect',
127 choices=sq.supported_formats(),
128 help='Assume input files are of given FORMAT. Choices: %(choices)s. '
129 'Default: %(default)s.')
131 p.add_argument(
132 '--kind', '-k',
133 type=csvtype(sq.supported_content_kinds()),
134 dest='kinds',
135 help='Restrict meta-data scanning to given content kinds. '
136 'KINDS is a comma-separated list of content kinds, choices: %s. '
137 'By default, all content kinds are indexed.'
138 % ', '.join(sq.supported_content_kinds()))
140 p.add_argument(
141 '--persistent', '-p',
142 dest='persistent',
143 metavar='NAME',
144 help='Create/use persistent selection with given NAME. Persistent '
145 'selections can be used to speed up startup of Squirrel-based '
146 'applications.')
148 p.add_argument(
149 '--update', '-u',
150 dest='update',
151 action='store_true',
152 default=False,
153 help='Allow adding paths and datasets to existing persistent '
154 'selection.')
156 p.add_argument(
157 '--dataset', '-d',
158 dest='datasets',
159 default=[],
160 action='append',
161 metavar='FILE',
162 help='Add directories/files/sources from dataset description file. '
163 'This option can be repeated to add multiple datasets.')
166def squirrel_from_selection_arguments(args):
167 from pyrocko.squirrel import base, dataset
169 datasets = [
170 dataset.read_dataset(dataset_path) for dataset_path in args.datasets]
172 persistents = [ds.persistent or '' for ds in datasets if ds.persistent]
173 if args.persistent:
174 persistent = args.persistent
175 elif persistents:
176 persistent = persistents[0]
177 if not all(p == persistents for p in persistents[1:]):
178 raise error.SquirrelError(
179 'Given datasets specify different `persistent` settings.')
181 if persistent:
182 logger.info(
183 'Persistent selection requested by dataset: %s' % persistent)
184 else:
185 persistent = None
187 else:
188 persistent = None
190 squirrel = base.Squirrel(persistent=persistent)
192 if persistent and not squirrel.is_new():
193 if not args.update:
194 logger.info(
195 'Using existing persistent selection: %s' % persistent)
196 if args.paths or datasets:
197 logger.info(
198 'Avoiding dataset rescan. Use --update/-u to '
199 'rescan or add items to existing persistent selection.')
201 return squirrel
203 else:
204 logger.info(
205 'Updating existing persistent selection: %s' % persistent)
207 kinds = args.kinds or None
208 if args.paths:
209 squirrel.add(
210 args.paths,
211 check=args.check,
212 format=args.format,
213 kinds=kinds,
214 include=args.include,
215 exclude=args.exclude)
217 for ds in datasets:
218 squirrel.add_dataset(ds, check=args.check)
220 return squirrel
223def add_query_arguments(p, without=[]):
224 if 'codes' not in without:
225 p.add_argument(
226 '--codes',
227 dest='codes',
228 metavar='CODES',
229 help='Code pattern to query (STA, NET.STA, NET.STA.LOC, '
230 'NET.STA.LOC.CHA, NET.STA.LOC.CHA.EXTRA, '
231 'AGENCY.NET.STA.LOC.CHA.EXTRA).')
233 if 'tmin' not in without:
234 p.add_argument(
235 '--tmin',
236 dest='tmin',
237 metavar='TIME',
238 help='Begin of time interval to query.')
240 if 'tmax' not in without:
241 p.add_argument(
242 '--tmax',
243 dest='tmax',
244 metavar='TIME',
245 help='End of time interval to query.')
247 if 'time' not in without:
248 p.add_argument(
249 '--time',
250 dest='time',
251 metavar='TIME',
252 help='Time instant to query.')
255def squirrel_query_from_arguments(args):
256 d = {}
258 if 'kinds' in args and args.kinds:
259 d['kind'] = args.kinds
260 if 'tmin' in args and args.tmin:
261 d['tmin'] = util.str_to_time_fillup(args.tmin)
262 if 'tmax' in args and args.tmax:
263 d['tmax'] = util.str_to_time_fillup(args.tmax)
264 if 'time' in args and args.time:
265 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
266 if 'codes' in args and args.codes:
267 d['codes'] = tuple(args.codes.split('.'))
269 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
270 raise error.SquirrelError(
271 'Options --tmin/--tmax and --time are mutually exclusive.')
272 return d
275class SquirrelCommand(object):
277 def add_parser(self, subparsers, *args, **kwargs):
278 return add_parser(subparsers, *args, **kwargs)
280 def add_selection_arguments(self, p):
281 return add_selection_arguments(p)
283 def add_query_arguments(self, p):
284 return add_query_arguments(p)
286 def squirrel_query_from_arguments(self, args):
287 return squirrel_query_from_arguments(args)
289 def squirrel_from_selection_arguments(self, args):
290 return squirrel_from_selection_arguments(args)
292 def fail(self, message):
293 raise error.ToolError(message)
296__all__ = ['SquirrelCommand']