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 help='Set logger level. Choices: critical, error, warning, '
71 'info [default], debug.')
73 p.add_argument(
74 '--progress',
75 choices=['terminal', 'log', 'off'],
76 default='terminal',
77 help='Set how progress status is reported. Choices: terminal '
78 '[default], off.')
81def process_standard_arguments(parser, args):
82 loglevel = args.__dict__.pop('loglevel')
83 util.setup_logging(parser.prog, loglevel)
85 pmode = args.__dict__.pop('progress')
86 progress.set_default_viewer(pmode)
89def add_selection_arguments(p):
90 from pyrocko import squirrel as sq
92 p.add_argument(
93 'paths',
94 nargs='*',
95 help='Files and directories with waveforms, metadata and events.')
97 p.add_argument(
98 '--include',
99 dest='include',
100 metavar='REGEX',
101 help='If given, files are only included if their paths match the '
102 'regular expression PATTERN.')
104 p.add_argument(
105 '--exclude',
106 dest='exclude',
107 metavar='REGEX',
108 help='If given, files are only included if their paths do not match '
109 'the regular expression PATTERN.')
111 p.add_argument(
112 '--optimistic', '-o',
113 action='store_false',
114 dest='check',
115 default=True,
116 help='Disable checking file modification times.')
118 p.add_argument(
119 '--format', '-f',
120 dest='format',
121 metavar='FORMAT',
122 default='detect',
123 choices=sq.supported_formats(),
124 help='Assume input files are of given FORMAT. Choices: %(choices)s. '
125 'Default: %(default)s.')
127 p.add_argument(
128 '--kind', '-k',
129 type=csvtype(sq.supported_content_kinds()),
130 dest='kinds',
131 help='Restrict meta-data scanning to given content kinds. '
132 'KINDS is a comma-separated list of content kinds, choices: %s. '
133 'By default, all content kinds are indexed.'
134 % ', '.join(sq.supported_content_kinds()))
136 p.add_argument(
137 '--persistent', '-p',
138 dest='persistent',
139 metavar='NAME',
140 help='Create/use persistent selection with given NAME. Persistent '
141 'selections can be used to speed up startup of Squirrel-based '
142 'applications.')
144 p.add_argument(
145 '--update', '-u',
146 dest='update',
147 action='store_true',
148 default=False,
149 help='Allow adding paths and datasets to existing persistent '
150 'selection.')
152 p.add_argument(
153 '--dataset', '-d',
154 dest='datasets',
155 default=[],
156 action='append',
157 metavar='FILE',
158 help='Add directories/files/sources from dataset description file. '
159 'This option can be repeated to add multiple datasets.')
162def squirrel_from_selection_arguments(args):
163 from pyrocko.squirrel import base, dataset
165 datasets = [
166 dataset.read_dataset(dataset_path) for dataset_path in args.datasets]
168 persistents = [ds.persistent or '' for ds in datasets if ds.persistent]
169 if args.persistent:
170 persistent = args.persistent
171 elif persistents:
172 persistent = persistents[0]
173 if not all(p == persistents for p in persistents[1:]):
174 raise error.SquirrelError(
175 'Given datasets specify different `persistent` settings.')
177 if persistent:
178 logger.info(
179 'Persistent selection requested by dataset: %s' % persistent)
180 else:
181 persistent = None
183 else:
184 persistent = None
186 squirrel = base.Squirrel(persistent=persistent)
188 if persistent and not squirrel.is_new():
189 if not args.update:
190 logger.info(
191 'Using existing persistent selection: %s' % persistent)
192 if args.paths or datasets:
193 logger.info(
194 'Avoiding dataset rescan. Use --update/-u to '
195 'rescan or add items to existing persistent selection.')
197 return squirrel
199 else:
200 logger.info(
201 'Updating existing persistent selection: %s' % persistent)
203 kinds = args.kinds or None
204 if args.paths:
205 squirrel.add(
206 args.paths,
207 check=args.check,
208 format=args.format,
209 kinds=kinds,
210 include=args.include,
211 exclude=args.exclude)
213 for ds in datasets:
214 squirrel.add_dataset(ds, check=args.check)
216 return squirrel
219def add_query_arguments(p, without=[]):
220 if 'codes' not in without:
221 p.add_argument(
222 '--codes',
223 dest='codes',
224 metavar='CODES',
225 help='Code pattern to query (STA, NET.STA, NET.STA.LOC, '
226 'NET.STA.LOC.CHA, NET.STA.LOC.CHA.EXTRA, '
227 'AGENCY.NET.STA.LOC.CHA.EXTRA).')
229 if 'tmin' not in without:
230 p.add_argument(
231 '--tmin',
232 dest='tmin',
233 metavar='TIME',
234 help='Begin of time interval to query.')
236 if 'tmax' not in without:
237 p.add_argument(
238 '--tmax',
239 dest='tmax',
240 metavar='TIME',
241 help='End of time interval to query.')
243 if 'time' not in without:
244 p.add_argument(
245 '--time',
246 dest='time',
247 metavar='TIME',
248 help='Time instant to query.')
251def squirrel_query_from_arguments(args):
252 d = {}
254 if 'kinds' in args and args.kinds:
255 d['kind'] = args.kinds
256 if 'tmin' in args and args.tmin:
257 d['tmin'] = util.str_to_time_fillup(args.tmin)
258 if 'tmax' in args and args.tmax:
259 d['tmax'] = util.str_to_time_fillup(args.tmax)
260 if 'time' in args and args.time:
261 d['tmin'] = d['tmax'] = util.str_to_time_fillup(args.time)
262 if 'codes' in args and args.codes:
263 d['codes'] = tuple(args.codes.split('.'))
265 if ('tmin' in d and 'time' in d) or ('tmax' in d and 'time' in d):
266 raise error.SquirrelError(
267 'Options --tmin/--tmax and --time are mutually exclusive.')
268 return d
271class SquirrelCommand(object):
273 def add_parser(self, subparsers, *args, **kwargs):
274 return add_parser(subparsers, *args, **kwargs)
276 def add_selection_arguments(self, p):
277 return add_selection_arguments(p)
279 def add_query_arguments(self, p):
280 return add_query_arguments(p)
282 def squirrel_query_from_arguments(self, args):
283 return squirrel_query_from_arguments(args)
285 def squirrel_from_selection_arguments(self, args):
286 return squirrel_from_selection_arguments(args)
288 def fail(self, message):
289 raise error.ToolError(message)
292__all__ = ['SquirrelCommand']