Coverage for /opt/anaconda/lib/python3.7/site-packages/llama/flags/__init__.py : 75%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# (c) Stefan Countryman 2019
3"""
4A mixin for keeping track of the state of flags applied to an event, e.g.
5whether an event is in 'test' or 'observation' mode.
6"""
8import os
9import json
10import functools
11from collections import namedtuple
12from llama.classes import ImmutableDict
13from llama.utils import COLOR
16FLAG_TABLE_COLWIDTH = 20
17ALLOWED_VALUES = ImmutableDict({
18 "ROLE": ("test", "observation"),
19 "ONLINE": ("true", "false"),
20 "VETOED": ("true", "false"),
21 "UPLOAD": ("true", "false"),
22 "ICECUBE_UPLOAD": ("true", "false"),
23 "BLINDED_NEUTRINOS": ("true", "false"),
24 "MANUAL": ("true", "false"),
25})
28class FlagPreset(ImmutableDict):
29 """
30 A set of key, value pairs for ``llama`` flags to serve common use cases.
31 Adds an optional extra ``description`` field to ``ImmutableDict`` that can
32 be used to describe the intended use case for a ``FlagPreset``.
33 """
35 def __new__(cls, *args, description=None, **kwargs):
36 """
37 Specify a dictionary or iterable (same as ``ImmutableDict``) with which
38 to construct a ``FlagPreset``. Optionally provide a ``description`` of
39 which use case this ``FlagPreset`` is meant to serve.
40 """
41 new = ImmutableDict.__new__(cls, *args, **kwargs)
42 setattr(new, 'description', description)
43 return new
46FLAG_PRESETS = namedtuple(
47 "FlagPresets",
48 (
49 "DEFAULT",
50 "TRIGGERED_INTERNAL",
51 "TRIGGERED_PUBLIC",
52 "TRIGGERED_TEST",
53 "RERUN",
54 "MANUAL",
55 )
56)(
57 DEFAULT=FlagPreset(
58 {
59 "ROLE": "test",
60 "ONLINE": "true",
61 "VETOED": "false",
62 "UPLOAD": "false",
63 "ICECUBE_UPLOAD": "false",
64 "BLINDED_NEUTRINOS": "false",
65 "MANUAL": "false",
66 },
67 description=(
68 "Safe defaults that will use real neutrino data but won't send "
69 "slack alerts to anybody. Marks this event as a test event, which "
70 "will automatically veto pipeline steps that have potentially "
71 "dangerous side-effects."
72 ),
73 ),
74 TRIGGERED_INTERNAL=FlagPreset(
75 {
76 "ROLE": "observation",
77 "ONLINE": "true",
78 "VETOED": "false",
79 "UPLOAD": "true",
80 "ICECUBE_UPLOAD": "false",
81 "BLINDED_NEUTRINOS": "false",
82 "MANUAL": "false",
83 },
84 description=(
85 "Flag preset used for externally-triggered events that should be "
86 "kept within the LLAMA group. Will run the pipeline just like it "
87 "runs for real events, EXCEPT that IceCube and other public "
88 "partners will not be alerted. Results are uploaded to LLAMA team "
89 "Slack."
90 ),
91 ),
92 TRIGGERED_PUBLIC=FlagPreset(
93 {
94 "ROLE": "observation",
95 "ONLINE": "true",
96 "VETOED": "false",
97 "UPLOAD": "true",
98 "ICECUBE_UPLOAD": "true",
99 "BLINDED_NEUTRINOS": "false",
100 "MANUAL": "false",
101 },
102 description=(
103 "Flags used by real events. The full pipeline runs, and results "
104 "are automatically sent to LLAMA and IceCube teams."
105 ),
106 ),
107 TRIGGERED_TEST=FlagPreset(
108 {
109 "ROLE": "test",
110 "ONLINE": "true",
111 "VETOED": "false",
112 "UPLOAD": "false",
113 "ICECUBE_UPLOAD": "false",
114 "BLINDED_NEUTRINOS": "true",
115 "MANUAL": "false",
116 },
117 description=(
118 "Flags for test events sent by alert services like GCN or "
119 "LVAlert. Will use scrambled neutrino data. Use these if you want "
120 "to test the pipeline's ability to "
121 "react to a new alert without running expensive code or alerting "
122 "people."
123 ),
124 ),
125 RERUN=FlagPreset(
126 {
127 "ROLE": "observation",
128 "ONLINE": "true",
129 "VETOED": "false",
130 "UPLOAD": "false",
131 "ICECUBE_UPLOAD": "false",
132 "BLINDED_NEUTRINOS": "false",
133 "MANUAL": "false",
134 },
135 description=(
136 "Flags for rerunning an analysis as if it were real but without "
137 "uploading files to other people. Use this if you want to re-run "
138 "an analysis but don't want to upload results anywhere."
139 ),
140 ),
141 MANUAL=FlagPreset(
142 {
143 "ROLE": "observation",
144 "ONLINE": "true",
145 "VETOED": "false",
146 "UPLOAD": "true",
147 "ICECUBE_UPLOAD": "true",
148 "BLINDED_NEUTRINOS": "false",
149 "MANUAL": "true",
150 },
151 description=(
152 "Sets the pipeline to use real data and allow all possible file "
153 "uploads and pipeline steps, but prevents the LLAMA damon from "
154 "automatically updating the event by setting the MANUAL flag to "
155 "true. You can still manually generate files."
156 ),
157 ),
158)
161# all flag presets are allowed values
162assert all(all(val in ALLOWED_VALUES[flag]
163 for flag, val in getattr(FLAG_PRESETS, preset).items())
164 for preset in FLAG_PRESETS._fields)
165# all flag presets specify all required values defined in ALLOWED_VALUES
166assert all(set(list(getattr(FLAG_PRESETS, preset))) ==
167 set(list(ALLOWED_VALUES))
168 for preset in FLAG_PRESETS._fields)
171class FlagDict:
172 """
173 A dict-like interface to flags that sets and gets values from an on-disk
174 file containing flags for a specific event directory. *Note that values and
175 keys must both be strings.* For flags in ``FlagDict.ALLOWED_VALUES``,
176 the provided value can only be set to one of the allowed values. Either of
177 these dictionaries can be extended as necessary to provide extra defaults
178 and restrictions across all FlagDict instances.
180 Parameters
181 ----------
182 eventdir : str
183 The path to the event directory that contains these flags.
184 """
186 def __init__(self, eventdir):
187 self.eventdir = os.path.abspath(eventdir)
189 def __eq__(self, other):
190 if not hasattr(other, 'items'):
191 return False
192 return frozenset(self.items()) == frozenset(other.items())
194 # default flags. since this is the global variable for flags, feel free to
195 # extend it with new flag defaults.
196 DEFAULT_FLAGS = FLAG_PRESETS.DEFAULT
198 # preset options for flags. use these for common use cases.
199 PRESETS = FLAG_PRESETS
201 # some flags only have certain allowed values.
202 ALLOWED_VALUES = ALLOWED_VALUES
204 assert set(list(ALLOWED_VALUES)) == set(list(DEFAULT_FLAGS))
206 def _read(self):
207 """Read in the JSON flag file as a dict."""
208 flags = dict(self.DEFAULT_FLAGS)
209 try:
210 with open(self._flagfile) as flagfile:
211 flags.update(json.load(flagfile))
212 except IOError:
213 pass
214 return flags
216 def _write(self, flags):
217 """Overwrite the JSON flag file with the provided dict."""
218 with open(self._flagfile, 'w') as flagfile:
219 json.dump(flags, flagfile, indent=4, sort_keys=True)
221 @property
222 def _flagfile(self):
223 """The file storing the flags for this event."""
224 return os.path.join(self.eventdir, "FLAGS.json")
226 def __contains__(self, item):
227 return item in self.keys()
229 def items(self):
230 """D.items() -> list of D's (key, value) pairs, as 2-tuples."""
231 return self._read().items()
233 def values(self):
234 """D.values() -> list of D's values"""
235 return self._read().values()
237 def keys(self):
238 """D.keys() -> list of D's keys"""
239 return self._read().keys()
241 def __getitem__(self, key):
242 # because namedtuples are stupid and use '__getitem__' to get their
243 # named properties, you need to check the fields manually to avoid
244 # reading in the file path when recovering the eventdir.
245 # for i, item in enumerate(self):
246 # if self._fields == key:
247 # return item
248 return self._read()[key]
250 def __setitem__(self, key, value):
251 if not (str(value) == value and str(key) == key):
252 raise ValueError("Can only set flags equal to strings.")
253 if key in self.ALLOWED_VALUES:
254 if value not in self.ALLOWED_VALUES[key]:
255 raise ValueError(("Invalid value {} for flag {}; must be one "
256 "of: {}. See ``FlagDict.ALLOWED_VALUES``"
257 "").format(value, key,
258 self.ALLOWED_VALUES[key]))
259 dic = self._read()
260 # TODO put some kind of lock here to prevent race condition corruption
261 dic[key] = str(value)
262 self._write(dic)
264 def __delitem__(self, key):
265 dic = self._read()
266 del dic[key]
267 self._write(dic)
269 def __iter__(self):
270 """Iterate over just the keys, mimicking ``dict`` behavior.
271 Implementation does just yields the keys from the preloaded list of
272 keys, so it's not much of an iterator."""
273 for k in self.keys():
274 yield k
276 def __repr__(self):
277 return repr(self._read())
279 def __str__(self):
280 return str(self._read())
282 def update(self, *E, **F):
283 """(Same as ``dict.update``.) ``self.update([E, ]**F)`` -> None.
284 Update ``self`` from dict/iterable E and F.
286 If E is present and has a .keys() method, then does:
287 for k in E: self[k] = E[k]
288 If E is present and lacks a .keys() method, then does:
289 for k, v in E: self[k] = v
290 In either case, this is followed by:
291 for k in F: self[k] = F[k]
292 """
293 if E:
294 if len(E) != 1:
295 raise ValueError(("Can only pass a single dict-like object to "
296 "``update``. Got: {}").format(E))
297 dictlike = E[0]
298 if hasattr(dictlike, 'keys'):
299 for k in dictlike.keys():
300 self[k] = dictlike[k]
301 else:
302 for k, v in dictlike:
303 self[k] = v
304 for k in F:
305 self[k] = F[k]
308def flag_table(flagdict, color=True, emphasize=()):
309 """
310 Print flags into a nice table for terminals. If ``color`` is ``True``,
311 add terminal color escape codes for emphasis. Specify flags that should be
312 emphasized by listing the flag names in ``emphasize`` (does not apply if
313 ``color`` is ``False``).
314 """
315 delim = '|'
316 col_hline = '-'*(2+FLAG_TABLE_COLWIDTH)
317 cfmt = "{0:<{1}}".format
318 if color:
319 delim = COLOR.CLEAR + delim + COLOR.CLEAR
320 hline = (COLOR.CLEAR + '\n' + COLOR.CLEAR + col_hline + '+' +
321 col_hline + COLOR.CLEAR + '\n')
322 head = COLOR.RED + COLOR.BOLD
323 emph = lambda k, v: (
324 (
325 COLOR.BOLD+COLOR.BLUE,
326 cfmt(
327 COLOR.UNDERLINE+k+COLOR.CLEAR,
328 FLAG_TABLE_COLWIDTH+len(COLOR.UNDERLINE+COLOR.CLEAR)
329 ),
330 COLOR.BOLD+COLOR.YELLOW,
331 cfmt(
332 COLOR.UNDERLINE+v+COLOR.CLEAR,
333 FLAG_TABLE_COLWIDTH+len(COLOR.UNDERLINE+COLOR.CLEAR)
334 ),
335 ) if k in emphasize else (
336 COLOR.BLUE,
337 cfmt(k, FLAG_TABLE_COLWIDTH),
338 COLOR.GREEN,
339 cfmt(v, FLAG_TABLE_COLWIDTH)
340 )
341 )
342 else:
343 head = ''
344 emph = lambda k, v: ('', cfmt(k, FLAG_TABLE_COLWIDTH),
345 '', cfmt(v, FLAG_TABLE_COLWIDTH))
346 hline = ('\n'+col_hline+'+'+col_hline+'\n')
347 fmt = (" {}{}" + COLOR.CLEAR + " " + delim + " {}{}" + COLOR.CLEAR)
348 rows = [fmt.format(*emph(k, v)) for k, v in sorted(flagdict.items())]
349 header = (fmt.format(head, cfmt("FLAG NAME", FLAG_TABLE_COLWIDTH),
350 head, cfmt("VALUE", FLAG_TABLE_COLWIDTH)) +
351 hline)
352 # hline.replace('-', '='))
353 # return (header + hline.join(rows))
354 return header + '\n'.join(rows)
357class FlagsMixin:
358 """
359 A mixin that can read from and write to flags the 'FLAGS.json' file of
360 any object with an ``eventdir`` property.
361 """
363 @property
364 def flags(self):
365 """A list of flags that apply to this instance, e.g. whether this is a
366 testing or production event. See ``llama.flags.FlagDict`` for details.
368 Returns
369 -------
370 flags : llama.flags.FlagDict
371 A collection of flags for this event directory with a
372 dictionary-like interface. Flags are read from file using standard
373 dictionary syntax, e.g. ``self.flags['foo']``, with default flags
374 added from ``FlagDict.DEFAULT_FLAGS``. Setting new flags with
375 something like ``self.flags['foo'] = bar`` will write the new flags
376 to file. Changes to flags are not automatically version-controlled
377 because they are committed as part of event directory state prior
378 to any file generation attempts.
379 """
380 return FlagDict(self.eventdir)
382 @flags.setter
383 def flags(self, flags):
384 flaghandler = self.flags
385 for oldflag in set(flaghandler) - set(flags):
386 del flaghandler[oldflag]
387 for newflag, value in flags.items():
388 flaghandler[newflag] = value
390 def decorate_checkout(func):
391 """
392 Copy flags to the temporary directory on checkout (since they can
393 control execution behavior).
394 """
396 @functools.wraps(func)
397 def wrapper(self, *args, **kwargs):
398 """
399 Copy flags to the temporary directory on checkout (since they can
400 control execution behavior).
401 """
402 tmp_self = func(self, *args, **kwargs)
403 tmp_self.flags = self.flags
404 return tmp_self
406 return wrapper