Merge pull request #209 from yilei/push_up_to_497985387

Push up to 497985387

CHANGELOG.md

@@ -7,6 +7,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com).
 ## Unreleased
 
 ### Changed
+*   (testing) Added `@flagsaver.as_parsed`: this allows saving/restoring flags
+    using string values as if parsed from the command line and will also reflect
+    other flag states after command line parsing, e.g. `.present` is set.
 
 *   (logging) If no log dir is specified `logging.find_log_dir()` now falls back
     to `tempfile.gettempdir()` instead of `/tmp/`.

absl/flags/_flagvalues.py

@@ -792,8 +792,10 @@ def get_value():
           continue
 
       if flag is not None:
+        # LINT.IfChange
         flag.parse(value)
         flag.using_default_value = False
+        # LINT.ThenChange(../testing/flagsaver.py:flag_override_parsing)
       else:
         unparsed_names_and_args.append((name, arg))
 

absl/testing/BUILD

@@ -212,6 +212,7 @@ py_test(
     deps = [
         ":absltest",
         ":flagsaver",
+        ":parameterized",
         "//absl/flags",
     ],
 )

absl/testing/flagsaver.py

@@ -11,6 +11,7 @@
 # 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
@@ -49,6 +50,36 @@ def some_func():
     finally:
       flagsaver.restore_flag_values(saved_flag_values)
 
+    # Use the parsing version to emulate users providing the flags.
+    # Note that all flags must be provided as strings (unparsed).
+    @flagsaver.as_parsed(some_int_flag='123')
+    def some_func():
+      # Because the flag was parsed it is considered "present".
+      assert FLAGS.some_int_flag.present
+      do_stuff()
+
+    # flagsaver.as_parsed() can also be used as a context manager just like
+    # flagsaver.flagsaver()
+    with flagsaver.as_parsed(some_int_flag='123'):
+      do_stuff()
+
+    # The flagsaver.as_parsed() interface also supports FlagHolder objects.
+    @flagsaver.as_parsed((module.FOO_FLAG, 'foo'), (other_mod.BAR_FLAG, '23'))
+    def some_func():
+      do_stuff()
+
+    # Using as_parsed with a multi_X flag requires a sequence of strings.
+    @flagsaver.as_parsed(some_multi_int_flag=['123', '456'])
+    def some_func():
+      assert FLAGS.some_multi_int_flag.present
+      do_stuff()
+
+    # If a flag name includes non-identifier characters it can be specified like
+    # so:
+    @flagsaver.as_parsed(**{'i-like-dashes': 'true'})
+    def some_func():
+      do_stuff()
+
 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.
@@ -58,14 +89,16 @@ def some_func():
 and then restore flag values, the added flag will be deleted with no errors.
 """
 
+import collections
 import functools
 import inspect
-from typing import overload, Any, Callable, Mapping, Tuple, TypeVar
+from typing import overload, Any, Callable, Mapping, Tuple, TypeVar, Type, Sequence, Union
 
 from absl import flags
 
 FLAGS = flags.FLAGS
 
+
 # The type of pre/post wrapped functions.
 _CallableT = TypeVar('_CallableT', bound=Callable)
 
@@ -83,8 +116,86 @@ def flagsaver(func: _CallableT) -> _CallableT:
 
 def flagsaver(*args, **kwargs):
   """The main flagsaver interface. See module doc for usage."""
+  return _construct_overrider(_FlagOverrider, *args, **kwargs)
+
+
+@overload
+def as_parsed(*args: Tuple[flags.FlagHolder, Union[str, Sequence[str]]],
+              **kwargs: Union[str, Sequence[str]]) -> '_ParsingFlagOverrider':
+  ...
+
+
+@overload
+def as_parsed(func: _CallableT) -> _CallableT:
+  ...
+
+
+def as_parsed(*args, **kwargs):
+  """Overrides flags by parsing strings, saves flag state similar to flagsaver.
+
+  This function can be used as either a decorator or context manager similar to
+  flagsaver.flagsaver(). However, where flagsaver.flagsaver() directly sets the
+  flags to new values, this function will parse the provided arguments as if
+  they were provided on the command line. Among other things, this will cause
+  `FLAGS['flag_name'].parsed == True`.
+
+  A note on unparsed input: For many flag types, the unparsed version will be
+  a single string. However for multi_x (multi_string, multi_integer, multi_enum)
+  the unparsed version will be a Sequence of strings.
+
+  Args:
+    *args: Tuples of FlagHolders and their unparsed value.
+    **kwargs: The keyword args are flag names, and the values are unparsed
+      values.
+
+  Returns:
+    _ParsingFlagOverrider that serves as a context manager or decorator. Will
+    save previous flag state and parse new flags, then on cleanup it will
+    restore the previous flag state.
+  """
+  return _construct_overrider(_ParsingFlagOverrider, *args, **kwargs)
+
+
+# NOTE: the order of these overload declarations matters. The type checker will
+# pick the first match which could be incorrect.
+@overload
+def _construct_overrider(
+    flag_overrider_cls: Type['_ParsingFlagOverrider'],
+    *args: Tuple[flags.FlagHolder, Union[str, Sequence[str]]],
+    **kwargs: Union[str, Sequence[str]]) -> '_ParsingFlagOverrider':
+  ...
+
+
+@overload
+def _construct_overrider(flag_overrider_cls: Type['_FlagOverrider'],
+                         *args: Tuple[flags.FlagHolder, Any],
+                         **kwargs: Any) -> '_FlagOverrider':
+  ...
+
+
+@overload
+def _construct_overrider(flag_overrider_cls: Type['_FlagOverrider'],
+                         func: _CallableT) -> _CallableT:
+  ...
+
+
+def _construct_overrider(flag_overrider_cls, *args, **kwargs):
+  """Handles the args/kwargs returning an instance of flag_overrider_cls.
+
+  If flag_overrider_cls is _FlagOverrider then values should be native python
+  types matching the python types. Otherwise if flag_overrider_cls is
+  _ParsingFlagOverrider the values should be strings or sequences of strings.
+
+  Args:
+    flag_overrider_cls: The class that will do the overriding.
+    *args: Tuples of FlagHolder and the new flag value.
+    **kwargs: Keword args mapping flag name to new flag value.
+
+  Returns:
+    A _FlagOverrider to be used as a decorator or context manager.
+  """
   if not args:
-    return _FlagOverrider(**kwargs)
+    return flag_overrider_cls(**kwargs)
   # args can be [func] if used as `@flagsaver` instead of `@flagsaver(...)`
   if len(args) == 1 and callable(args[0]):
     if kwargs:
@@ -93,7 +204,7 @@ def flagsaver(*args, **kwargs):
     func = args[0]
     if inspect.isclass(func):
       raise TypeError('@flagsaver.flagsaver cannot be applied to a class.')
-    return _wrap(func, {})
+    return _wrap(flag_overrider_cls, func, {})
   # args can be a list of (FlagHolder, value) pairs.
   # In which case they augment any specified kwargs.
   for arg in args:
@@ -105,7 +216,7 @@ def flagsaver(*args, **kwargs):
     if holder.name in kwargs:
       raise ValueError('Cannot set --%s multiple times' % holder.name)
     kwargs[holder.name] = value
-  return _FlagOverrider(**kwargs)
+  return flag_overrider_cls(**kwargs)
 
 
 def save_flag_values(
@@ -144,13 +255,27 @@ def restore_flag_values(saved_flag_values: Mapping[str, Mapping[str, Any]],
       flag_values[name].__dict__ = saved
 
 
-def _wrap(func: _CallableT, overrides: Mapping[str, Any]) -> _CallableT:
+@overload
+def _wrap(flag_overrider_cls: Type['_FlagOverrider'], func: _CallableT,
+          overrides: Mapping[str, Any]) -> _CallableT:
+  ...
+
+
+@overload
+def _wrap(flag_overrider_cls: Type['_ParsingFlagOverrider'], func: _CallableT,
+          overrides: Mapping[str, Union[str, Sequence[str]]]) -> _CallableT:
+  ...
+
+
+def _wrap(flag_overrider_cls, func, overrides):
   """Creates a wrapper function that saves/restores flag values.
 
   Args:
+    flag_overrider_cls: The class that will be used as a context manager.
     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.
+      saving the original flag state. The type of the values depends on if
+      _FlagOverrider or _ParsingFlagOverrider was specified.
 
   Returns:
     A wrapped version of func.
@@ -159,7 +284,7 @@ def _wrap(func: _CallableT, overrides: Mapping[str, Any]) -> _CallableT:
   @functools.wraps(func)
   def _flagsaver_wrapper(*args, **kwargs):
     """Wrapper function that saves and restores flags."""
-    with _FlagOverrider(**overrides):
+    with flag_overrider_cls(**overrides):
       return func(*args, **kwargs)
 
   return _flagsaver_wrapper
@@ -179,7 +304,7 @@ def __init__(self, **overrides: Any):
   def __call__(self, func: _CallableT) -> _CallableT:
     if inspect.isclass(func):
       raise TypeError('flagsaver cannot be applied to a class.')
-    return _wrap(func, self._overrides)
+    return _wrap(self.__class__, func, self._overrides)
 
   def __enter__(self):
     self._saved_flag_values = save_flag_values(FLAGS)
@@ -194,6 +319,55 @@ def __exit__(self, exc_type, exc_value, traceback):
     restore_flag_values(self._saved_flag_values, FLAGS)
 
 
+class _ParsingFlagOverrider(_FlagOverrider):
+  """Context manager for overriding flags.
+
+  Simulates command line parsing.
+
+  This is simlar to _FlagOverrider except that all **overrides should be
+  strings or sequences of strings, and when context is entered this class calls
+  .parse(value)
+
+  This results in the flags having .present set properly.
+  """
+
+  def __init__(self, **overrides: Union[str, Sequence[str]]):
+    for flag_name, new_value in overrides.items():
+      if isinstance(new_value, str):
+        continue
+      if (isinstance(new_value, collections.abc.Sequence) and
+          all(isinstance(single_value, str) for single_value in new_value)):
+        continue
+      raise TypeError(
+          f'flagsaver.as_parsed() cannot parse {flag_name}. Expected a single '
+          f'string or sequence of strings but {type(new_value)} was provided.')
+    super().__init__(**overrides)
+
+  def __enter__(self):
+    self._saved_flag_values = save_flag_values(FLAGS)
+    try:
+      for flag_name, unparsed_value in self._overrides.items():
+        # LINT.IfChange(flag_override_parsing)
+        FLAGS[flag_name].parse(unparsed_value)
+        FLAGS[flag_name].using_default_value = False
+        # LINT.ThenChange()
+
+      # Perform the validation on all modified flags. This is something that
+      # FLAGS._set_attributes() does for you in _FlagOverrider.
+      for flag_name in self._overrides:
+        FLAGS._assert_validators(FLAGS[flag_name].validators)
+
+    except KeyError as e:
+      # If a flag doesn't exist, an UnrecognizedFlagError is more specific.
+      restore_flag_values(self._saved_flag_values, FLAGS)
+      raise flags.UnrecognizedFlagError('Unknown command line flag.') from e
+
+    except:
+      # It may fail because of flag validators or general parsing issues.
+      restore_flag_values(self._saved_flag_values, FLAGS)
+      raise
+
+
 def _copy_flag_dict(flag: flags.Flag) -> Mapping[str, Any]:
   """Returns a copy of the flag object's ``__dict__``.