1# Copyright 2017 The Abseil Authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""Module to enforce different constraints on flags.
16
17Flags validators can be registered using following functions / decorators::
18
19    flags.register_validator
20    @flags.validator
21    flags.register_multi_flags_validator
22    @flags.multi_flags_validator
23
24Three convenience functions are also provided for common flag constraints::
25
26    flags.mark_flag_as_required
27    flags.mark_flags_as_required
28    flags.mark_flags_as_mutual_exclusive
29    flags.mark_bool_flags_as_mutual_exclusive
30
31See their docstring in this module for a usage manual.
32
33Do NOT import this module directly. Import the flags package and use the
34aliases defined at the package level instead.
35"""
36
37import warnings
38
39from absl.flags import _exceptions
40from absl.flags import _flagvalues
41from absl.flags import _validators_classes
42
43
44def register_validator(flag_name,
45                       checker,
46                       message='Flag validation failed',
47                       flag_values=_flagvalues.FLAGS):
48  """Adds a constraint, which will be enforced during program execution.
49
50  The constraint is validated when flags are initially parsed, and after each
51  change of the corresponding flag's value.
52
53  Args:
54    flag_name: str | FlagHolder, name or holder of the flag to be checked.
55        Positional-only parameter.
56    checker: callable, a function to validate the flag.
57
58        * input - A single positional argument: The value of the corresponding
59          flag (string, boolean, etc.  This value will be passed to checker
60          by the library).
61        * output - bool, True if validator constraint is satisfied.
62          If constraint is not satisfied, it should either ``return False`` or
63          ``raise flags.ValidationError(desired_error_message)``.
64
65    message: str, error text to be shown to the user if checker returns False.
66        If checker raises flags.ValidationError, message from the raised
67        error will be shown.
68    flag_values: flags.FlagValues, optional FlagValues instance to validate
69        against.
70
71  Raises:
72    AttributeError: Raised when flag_name is not registered as a valid flag
73        name.
74    ValueError: Raised when flag_values is non-default and does not match the
75        FlagValues of the provided FlagHolder instance.
76  """
77  flag_name, flag_values = _flagvalues.resolve_flag_ref(flag_name, flag_values)
78  v = _validators_classes.SingleFlagValidator(flag_name, checker, message)
79  _add_validator(flag_values, v)
80
81
82def validator(flag_name, message='Flag validation failed',
83              flag_values=_flagvalues.FLAGS):
84  """A function decorator for defining a flag validator.
85
86  Registers the decorated function as a validator for flag_name, e.g.::
87
88      @flags.validator('foo')
89      def _CheckFoo(foo):
90        ...
91
92  See :func:`register_validator` for the specification of checker function.
93
94  Args:
95    flag_name: str | FlagHolder, name or holder of the flag to be checked.
96        Positional-only parameter.
97    message: str, error text to be shown to the user if checker returns False.
98        If checker raises flags.ValidationError, message from the raised
99        error will be shown.
100    flag_values: flags.FlagValues, optional FlagValues instance to validate
101        against.
102  Returns:
103    A function decorator that registers its function argument as a validator.
104  Raises:
105    AttributeError: Raised when flag_name is not registered as a valid flag
106        name.
107  """
108
109  def decorate(function):
110    register_validator(flag_name, function,
111                       message=message,
112                       flag_values=flag_values)
113    return function
114  return decorate
115
116
117def register_multi_flags_validator(flag_names,
118                                   multi_flags_checker,
119                                   message='Flags validation failed',
120                                   flag_values=_flagvalues.FLAGS):
121  """Adds a constraint to multiple flags.
122
123  The constraint is validated when flags are initially parsed, and after each
124  change of the corresponding flag's value.
125
126  Args:
127    flag_names: [str | FlagHolder], a list of the flag names or holders to be
128        checked. Positional-only parameter.
129    multi_flags_checker: callable, a function to validate the flag.
130
131        * input - dict, with keys() being flag_names, and value for each key
132            being the value of the corresponding flag (string, boolean, etc).
133        * output - bool, True if validator constraint is satisfied.
134            If constraint is not satisfied, it should either return False or
135            raise flags.ValidationError.
136
137    message: str, error text to be shown to the user if checker returns False.
138        If checker raises flags.ValidationError, message from the raised
139        error will be shown.
140    flag_values: flags.FlagValues, optional FlagValues instance to validate
141        against.
142
143  Raises:
144    AttributeError: Raised when a flag is not registered as a valid flag name.
145    ValueError: Raised when multiple FlagValues are used in the same
146        invocation. This can occur when FlagHolders have different `_flagvalues`
147        or when str-type flag_names entries are present and the `flag_values`
148        argument does not match that of provided FlagHolder(s).
149  """
150  flag_names, flag_values = _flagvalues.resolve_flag_refs(
151      flag_names, flag_values)
152  v = _validators_classes.MultiFlagsValidator(
153      flag_names, multi_flags_checker, message)
154  _add_validator(flag_values, v)
155
156
157def multi_flags_validator(flag_names,
158                          message='Flag validation failed',
159                          flag_values=_flagvalues.FLAGS):
160  """A function decorator for defining a multi-flag validator.
161
162  Registers the decorated function as a validator for flag_names, e.g.::
163
164      @flags.multi_flags_validator(['foo', 'bar'])
165      def _CheckFooBar(flags_dict):
166        ...
167
168  See :func:`register_multi_flags_validator` for the specification of checker
169  function.
170
171  Args:
172    flag_names: [str | FlagHolder], a list of the flag names or holders to be
173        checked. Positional-only parameter.
174    message: str, error text to be shown to the user if checker returns False.
175        If checker raises flags.ValidationError, message from the raised
176        error will be shown.
177    flag_values: flags.FlagValues, optional FlagValues instance to validate
178        against.
179
180  Returns:
181    A function decorator that registers its function argument as a validator.
182
183  Raises:
184    AttributeError: Raised when a flag is not registered as a valid flag name.
185  """
186
187  def decorate(function):
188    register_multi_flags_validator(flag_names,
189                                   function,
190                                   message=message,
191                                   flag_values=flag_values)
192    return function
193
194  return decorate
195
196
197def mark_flag_as_required(flag_name, flag_values=_flagvalues.FLAGS):
198  """Ensures that flag is not None during program execution.
199
200  Registers a flag validator, which will follow usual validator rules.
201  Important note: validator will pass for any non-``None`` value, such as
202  ``False``, ``0`` (zero), ``''`` (empty string) and so on.
203
204  If your module might be imported by others, and you only wish to make the flag
205  required when the module is directly executed, call this method like this::
206
207      if __name__ == '__main__':
208        flags.mark_flag_as_required('your_flag_name')
209        app.run()
210
211  Args:
212    flag_name: str | FlagHolder, name or holder of the flag.
213        Positional-only parameter.
214    flag_values: flags.FlagValues, optional :class:`~absl.flags.FlagValues`
215        instance where the flag is defined.
216  Raises:
217    AttributeError: Raised when flag_name is not registered as a valid flag
218        name.
219    ValueError: Raised when flag_values is non-default and does not match the
220        FlagValues of the provided FlagHolder instance.
221  """
222  flag_name, flag_values = _flagvalues.resolve_flag_ref(flag_name, flag_values)
223  if flag_values[flag_name].default is not None:
224    warnings.warn(
225        'Flag --%s has a non-None default value; therefore, '
226        'mark_flag_as_required will pass even if flag is not specified in the '
227        'command line!' % flag_name,
228        stacklevel=2)
229  register_validator(
230      flag_name,
231      lambda value: value is not None,
232      message='Flag --{} must have a value other than None.'.format(flag_name),
233      flag_values=flag_values)
234
235
236def mark_flags_as_required(flag_names, flag_values=_flagvalues.FLAGS):
237  """Ensures that flags are not None during program execution.
238
239  If your module might be imported by others, and you only wish to make the flag
240  required when the module is directly executed, call this method like this::
241
242      if __name__ == '__main__':
243        flags.mark_flags_as_required(['flag1', 'flag2', 'flag3'])
244        app.run()
245
246  Args:
247    flag_names: Sequence[str | FlagHolder], names or holders of the flags.
248    flag_values: flags.FlagValues, optional FlagValues instance where the flags
249        are defined.
250  Raises:
251    AttributeError: If any of flag name has not already been defined as a flag.
252  """
253  for flag_name in flag_names:
254    mark_flag_as_required(flag_name, flag_values)
255
256
257def mark_flags_as_mutual_exclusive(flag_names, required=False,
258                                   flag_values=_flagvalues.FLAGS):
259  """Ensures that only one flag among flag_names is not None.
260
261  Important note: This validator checks if flag values are ``None``, and it does
262  not distinguish between default and explicit values. Therefore, this validator
263  does not make sense when applied to flags with default values other than None,
264  including other false values (e.g. ``False``, ``0``, ``''``, ``[]``). That
265  includes multi flags with a default value of ``[]`` instead of None.
266
267  Args:
268    flag_names: [str | FlagHolder], names or holders of flags.
269        Positional-only parameter.
270    required: bool. If true, exactly one of the flags must have a value other
271        than None. Otherwise, at most one of the flags can have a value other
272        than None, and it is valid for all of the flags to be None.
273    flag_values: flags.FlagValues, optional FlagValues instance where the flags
274        are defined.
275
276  Raises:
277    ValueError: Raised when multiple FlagValues are used in the same
278        invocation. This can occur when FlagHolders have different `_flagvalues`
279        or when str-type flag_names entries are present and the `flag_values`
280        argument does not match that of provided FlagHolder(s).
281  """
282  flag_names, flag_values = _flagvalues.resolve_flag_refs(
283      flag_names, flag_values)
284  for flag_name in flag_names:
285    if flag_values[flag_name].default is not None:
286      warnings.warn(
287          'Flag --{} has a non-None default value. That does not make sense '
288          'with mark_flags_as_mutual_exclusive, which checks whether the '
289          'listed flags have a value other than None.'.format(flag_name),
290          stacklevel=2)
291
292  def validate_mutual_exclusion(flags_dict):
293    flag_count = sum(1 for val in flags_dict.values() if val is not None)
294    if flag_count == 1 or (not required and flag_count == 0):
295      return True
296    raise _exceptions.ValidationError(
297        '{} one of ({}) must have a value other than None.'.format(
298            'Exactly' if required else 'At most', ', '.join(flag_names)))
299
300  register_multi_flags_validator(
301      flag_names, validate_mutual_exclusion, flag_values=flag_values)
302
303
304def mark_bool_flags_as_mutual_exclusive(flag_names, required=False,
305                                        flag_values=_flagvalues.FLAGS):
306  """Ensures that only one flag among flag_names is True.
307
308  Args:
309    flag_names: [str | FlagHolder], names or holders of flags.
310        Positional-only parameter.
311    required: bool. If true, exactly one flag must be True. Otherwise, at most
312        one flag can be True, and it is valid for all flags to be False.
313    flag_values: flags.FlagValues, optional FlagValues instance where the flags
314        are defined.
315
316  Raises:
317    ValueError: Raised when multiple FlagValues are used in the same
318        invocation. This can occur when FlagHolders have different `_flagvalues`
319        or when str-type flag_names entries are present and the `flag_values`
320        argument does not match that of provided FlagHolder(s).
321  """
322  flag_names, flag_values = _flagvalues.resolve_flag_refs(
323      flag_names, flag_values)
324  for flag_name in flag_names:
325    if not flag_values[flag_name].boolean:
326      raise _exceptions.ValidationError(
327          'Flag --{} is not Boolean, which is required for flags used in '
328          'mark_bool_flags_as_mutual_exclusive.'.format(flag_name))
329
330  def validate_boolean_mutual_exclusion(flags_dict):
331    flag_count = sum(bool(val) for val in flags_dict.values())
332    if flag_count == 1 or (not required and flag_count == 0):
333      return True
334    raise _exceptions.ValidationError(
335        '{} one of ({}) must be True.'.format(
336            'Exactly' if required else 'At most', ', '.join(flag_names)))
337
338  register_multi_flags_validator(
339      flag_names, validate_boolean_mutual_exclusion, flag_values=flag_values)
340
341
342def _add_validator(fv, validator_instance):
343  """Register new flags validator to be checked.
344
345  Args:
346    fv: flags.FlagValues, the FlagValues instance to add the validator.
347    validator_instance: validators.Validator, the validator to add.
348  Raises:
349    KeyError: Raised when validators work with a non-existing flag.
350  """
351  for flag_name in validator_instance.get_flags_names():
352    fv[flag_name].validators.append(validator_instance)
353