1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
|
# Copyright 2017 The Abseil Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Decorator and context manager for saving and restoring flag values.
There are many ways to save and restore. Always use the most convenient method
for a given use case.
Here are examples of each method. They all call ``do_stuff()`` while
``FLAGS.someflag`` is temporarily set to ``'foo'``::
from absl.testing import flagsaver
# Use a decorator which can optionally override flags via arguments.
@flagsaver.flagsaver(someflag='foo')
def some_func():
do_stuff()
# Use a decorator which can optionally override flags with flagholders.
@flagsaver.flagsaver((module.FOO_FLAG, 'foo'), (other_mod.BAR_FLAG, 23))
def some_func():
do_stuff()
# Use a decorator which does not override flags itself.
@flagsaver.flagsaver
def some_func():
FLAGS.someflag = 'foo'
do_stuff()
# Use a context manager which can optionally override flags via arguments.
with flagsaver.flagsaver(someflag='foo'):
do_stuff()
# Save and restore the flag values yourself.
saved_flag_values = flagsaver.save_flag_values()
try:
FLAGS.someflag = 'foo'
do_stuff()
finally:
flagsaver.restore_flag_values(saved_flag_values)
We save and restore a shallow copy of each Flag object's ``__dict__`` attribute.
This preserves all attributes of the flag, such as whether or not it was
overridden from its default value.
WARNING: Currently a flag that is saved and then deleted cannot be restored. An
exception will be raised. However if you *add* a flag after saving flag values,
and then restore flag values, the added flag will be deleted with no errors.
"""
import functools
import inspect
from typing import overload, Any, Callable, Mapping, Tuple, TypeVar
from absl import flags
FLAGS = flags.FLAGS
# The type of pre/post wrapped functions.
_CallableT = TypeVar('_CallableT', bound=Callable)
@overload
def flagsaver(*args: Tuple[flags.FlagHolder, Any],
**kwargs: Any) -> '_FlagOverrider':
...
@overload
def flagsaver(func: _CallableT) -> _CallableT:
...
def flagsaver(*args, **kwargs):
"""The main flagsaver interface. See module doc for usage."""
if not args:
return _FlagOverrider(**kwargs)
# args can be [func] if used as `@flagsaver` instead of `@flagsaver(...)`
if len(args) == 1 and callable(args[0]):
if kwargs:
raise ValueError(
"It's invalid to specify both positional and keyword parameters.")
func = args[0]
if inspect.isclass(func):
raise TypeError('@flagsaver.flagsaver cannot be applied to a class.')
return _wrap(func, {})
# args can be a list of (FlagHolder, value) pairs.
# In which case they augment any specified kwargs.
for arg in args:
if not isinstance(arg, tuple) or len(arg) != 2:
raise ValueError('Expected (FlagHolder, value) pair, found %r' % (arg,))
holder, value = arg
if not isinstance(holder, flags.FlagHolder):
raise ValueError('Expected (FlagHolder, value) pair, found %r' % (arg,))
if holder.name in kwargs:
raise ValueError('Cannot set --%s multiple times' % holder.name)
kwargs[holder.name] = value
return _FlagOverrider(**kwargs)
def save_flag_values(
flag_values: flags.FlagValues = FLAGS) -> Mapping[str, Mapping[str, Any]]:
"""Returns copy of flag values as a dict.
Args:
flag_values: FlagValues, the FlagValues instance with which the flag will be
saved. This should almost never need to be overridden.
Returns:
Dictionary mapping keys to values. Keys are flag names, values are
corresponding ``__dict__`` members. E.g. ``{'key': value_dict, ...}``.
"""
return {name: _copy_flag_dict(flag_values[name]) for name in flag_values}
def restore_flag_values(saved_flag_values: Mapping[str, Mapping[str, Any]],
flag_values: flags.FlagValues = FLAGS):
"""Restores flag values based on the dictionary of flag values.
Args:
saved_flag_values: {'flag_name': value_dict, ...}
flag_values: FlagValues, the FlagValues instance from which the flag will be
restored. This should almost never need to be overridden.
"""
new_flag_names = list(flag_values)
for name in new_flag_names:
saved = saved_flag_values.get(name)
if saved is None:
# If __dict__ was not saved delete "new" flag.
delattr(flag_values, name)
else:
if flag_values[name].value != saved['_value']:
flag_values[name].value = saved['_value'] # Ensure C++ value is set.
flag_values[name].__dict__ = saved
def _wrap(func: _CallableT, overrides: Mapping[str, Any]) -> _CallableT:
"""Creates a wrapper function that saves/restores flag values.
Args:
func: This will be called between saving flags and restoring flags.
overrides: Flag names mapped to their values. These flags will be set after
saving the original flag state.
Returns:
A wrapped version of func.
"""
@functools.wraps(func)
def _flagsaver_wrapper(*args, **kwargs):
"""Wrapper function that saves and restores flags."""
with _FlagOverrider(**overrides):
return func(*args, **kwargs)
return _flagsaver_wrapper
class _FlagOverrider(object):
"""Overrides flags for the duration of the decorated function call.
It also restores all original values of flags after decorated method
completes.
"""
def __init__(self, **overrides: Any):
self._overrides = overrides
self._saved_flag_values = None
def __call__(self, func: _CallableT) -> _CallableT:
if inspect.isclass(func):
raise TypeError('flagsaver cannot be applied to a class.')
return _wrap(func, self._overrides)
def __enter__(self):
self._saved_flag_values = save_flag_values(FLAGS)
try:
FLAGS._set_attributes(**self._overrides)
except:
# It may fail because of flag validators.
restore_flag_values(self._saved_flag_values, FLAGS)
raise
def __exit__(self, exc_type, exc_value, traceback):
restore_flag_values(self._saved_flag_values, FLAGS)
def _copy_flag_dict(flag: flags.Flag) -> Mapping[str, Any]:
"""Returns a copy of the flag object's ``__dict__``.
It's mostly a shallow copy of the ``__dict__``, except it also does a shallow
copy of the validator list.
Args:
flag: flags.Flag, the flag to copy.
Returns:
A copy of the flag object's ``__dict__``.
"""
copy = flag.__dict__.copy()
copy['_value'] = flag.value # Ensure correct restore for C++ flags.
copy['validators'] = list(flag.validators)
return copy
|
