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 if 'description' not in kwargs and 'help' in kwargs:
56 kwargs['description'] = kwargs['help']
57 p = subparsers.add_parser(*args, **kwargs)
58 add_standard_arguments(p)
59 return p
62def add_standard_arguments(p):
63 p.add_argument(
64 '--help', '-h',
65 action='help',
66 help='Show this help message and exit.')
68 p.add_argument(
69 '--loglevel',
70 choices=['critical', 'error', 'warning', 'info', 'debug'],
71 default='info',
72 metavar='LEVEL',
73 help='Set logger level. Choices: %(choices)s. Default: %(default)s.')
75 p.add_argument(
76 '--progress',
77 choices=['terminal', 'log', 'off'],
78 default='terminal',
79 metavar='DEST',
80 help='Set how progress status is reported. Choices: %(choices)s. '
81 'Default: %(default)s.')
84def process_standard_arguments(parser, args):
85 loglevel = args.__dict__.pop('loglevel')
86 util.setup_logging(parser.prog, loglevel)
88 pmode = args.__dict__.pop('progress')
89 progress.set_default_viewer(pmode)
92def add_selection_arguments(p):
93 from pyrocko import squirrel as sq
95 p.add_argument(
96 'paths',
97 nargs='*',
98 help='Files and directories with waveforms, metadata and events.')
100 p.add_argument(
101 '--include',
102 dest='include',
103 metavar='REGEX',
104 help='Only include files whose paths match REGEX. Examples: '
105 '--include=\'\\.MSEED$\' would only match files ending with '
106 '".MSEED". --include=\'\\.BH[EN]\\.\' would match paths '
107 'containing ".BHE." or ".BHN.". --include=\'/2011/\' would match '
108 'paths with a subdirectory "2011" in their path hierarchy.')
110 p.add_argument(
111 '--exclude',
112 dest='exclude',
113 metavar='REGEX',
114 help='If given, files are only included if their paths do not match '
115 'the regular expression pattern REGEX.')
117 p.add_argument(
118 '--optimistic', '-o',
119 action='store_false',
120 dest='check',
121 default=True,
122 help='Disable checking file modification times.')
124 p.add_argument(
125 '--format', '-f',
126 dest='format',
127 metavar='FORMAT',
128 default='detect',
129 choices=sq.supported_formats(),
130 help='Assume input files are of given FORMAT. Choices: %(choices)s. '
131 'Default: %(default)s.')
133 p.add_argument(
134 '--kind', '-k',
135 type=csvtype(sq.supported_content_kinds()),
136 dest='kinds',
137 help='Restrict meta-data scanning to given content kinds. '
138 'KINDS is a comma-separated list of content kinds, choices: %s. '
139 'By default, all content kinds are indexed.'
140 % ', '.join(sq.supported_content_kinds()))
142 p.add_argument(
143 '--persistent', '-p',
144 dest='persistent',
145 metavar='NAME',
146 help='Create/use persistent selection with given NAME. Persistent '
147 'selections can be used to speed up startup of Squirrel-based '
148 'applications.')
150 p.add_argument(
151 '--update', '-u',
152 dest='update',
153 action='store_true',
154 default=False,
155 help='Allow adding paths and datasets to existing persistent '
156 'selection.')
158 p.add_argument(
159 '--dataset', '-d',
160 dest='datasets',
161 default=[],
162 action='append',
163 metavar='FILE',
164 help='Add directories/files/sources from dataset description file. '
165 'This option can be repeated to add multiple datasets.')
168def squirrel_from_selection_arguments(args):
169 from pyrocko.squirrel import base, dataset
171 datasets = [
172 dataset.read_dataset(dataset_path) for dataset_path in args.datasets]
174 persistents = [ds.persistent or '' for ds in datasets if ds.persistent]
175 if args.persistent:
176 persistent = args.persistent
177 elif persistents:
178 persistent = persistents[0]
179 if not all(p == persistents for p in persistents[1:]):
180 raise error.SquirrelError(
181 'Given datasets specify different `persistent` settings.')
183 if persistent:
184 logger.info(
185 'Persistent selection requested by dataset: %s' % persistent)
186 else:
187 persistent = None
189 else:
190 persistent = None
192 squirrel = base.Squirrel(persistent=persistent)
194 if persistent and not squirrel.is_new():
195 if not args.update:
196 logger.info(
197 'Using existing persistent selection: %s' % persistent)
198 if args.paths or datasets:
199 logger.info(
200 'Avoiding dataset rescan. Use --update/-u to '
201 'rescan or add items to existing persistent selection.')
203 return squirrel
205 else:
206 logger.info(
207 'Updating existing persistent selection: %s' % persistent)
209 kinds = args.kinds or None
210 if args.paths:
211 squirrel.add(
212 args.paths,
213 check=args.check,
214 format=args.format,
215 kinds=kinds,
216 include=args.include,
217 exclude=args.exclude)
219 for ds in datasets:
220 squirrel.add_dataset(ds, check=args.check)
222 return squirrel
225def add_query_arguments(p, without=[]):
226 if 'codes' not in without:
227 p.add_argument(
228 '--codes',
229 dest='codes',
230 metavar='CODES',
231 help='Code pattern to query (STA, NET.STA, NET.STA.LOC, '
232 'NET.STA.LOC.CHA, NET.STA.LOC.CHA.EXTRA, '
233 'AGENCY.NET.STA.LOC.CHA.EXTRA).')
235 if 'tmin' not in without:
236 p.add_argument(
237 '--tmin',
238 dest='tmin',
239 metavar='TIME',
240 help='Begin of time interval to query.')
242 if 'tmax' not in without:
243 p.add_argument(
244 '--tmax',
245 dest='tmax',
246 metavar='TIME',
247 help='End of time interval to query.')
249 if 'time' not in without:
250 p.add_argument(
251 '--time',
252 dest='time',
253 metavar='TIME',
254 help='Time instant to query.')
257def squirrel_query_from_arguments(args):
258 d = {}
260 if 'kinds' in args and args.kinds:
261 d['kind'] = args.kinds
262 if 'tmin' in args and args.tmin:
263 d['tmin'] = util.str_to_time_fillup(args.tmin)
264 if 'tmax' in args and args.tmax:
265 d['tmax'] = util.str_to_time_fillup(args.tmax)
266 if 'time' in args and args.time:
267 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
268 if 'codes' in args and args.codes:
269 d['codes'] = tuple(args.codes.split('.'))
271 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
272 raise error.SquirrelError(
273 'Options --tmin/--tmax and --time are mutually exclusive.')
274 return d
277class SquirrelCommand(object):
279 def add_parser(self, subparsers, *args, **kwargs):
280 return add_parser(subparsers, *args, **kwargs)
282 def add_selection_arguments(self, p):
283 return add_selection_arguments(p)
285 def add_query_arguments(self, p, without=[]):
286 return add_query_arguments(p, without=without)
288 def squirrel_query_from_arguments(self, args):
289 return squirrel_query_from_arguments(args)
291 def squirrel_from_selection_arguments(self, args):
292 return squirrel_from_selection_arguments(args)
294 def fail(self, message):
295 raise error.ToolError(message)
297 def setup_subcommand(self, subparsers):
298 return self.add_parser(
299 subparsers, self.__class__.__name__, help='Undocumented.')
301 def setup(self, parser):
302 pass
304 def call(self, parser, args):
305 pass
308__all__ = ['SquirrelCommand']