Allow enum.Enum classes to be used directly with flags.

This adds a new flags define, `enum_class`, which is like `enum`, but uses
an `enum.Enum` class as the source of valid values. The resulting parsed
value is the corresponding Enum value.

Using enum_class is recommended over plain enum.

Python 2.7 note: The `enum` module is not required for absl in general, but
is required to use `enum_class`, or for absl tests to fully pass. If the
`enum` module isn't available, then absl flags can still be used, just not
the `enum_class` function. The enum34 package on PyPI is the suggested
package to use for enum support, and is the one absl assumes is used under
Python 2.7.

Python 2.7 and Bazel note: Because of how Bazel modifies PYTHONPATH, absl has
to do do some special sys.path import manipulation to make the enum module
work. In short, third party deps go before the stdlib on sys.path, so enum34's
enum module shadows Python 3's enum module, which causes problems because
enum34 is a subset of later Python 3's enum. So some sys.path futzing is done
to account for this under Bazel. This only happens if 'import enum' initially
fails, and only under Python 2.7.

PiperOrigin-RevId: 212301566

1 files changed, 51 insertions, 0 deletions

diff --git a/absl/flags/_argument_parser.py b/absl/flags/_argument_parser.py
index 15eda52..d391dc8 100644
--- a/absl/flags/_argument_parser.py
+++ b/absl/flags/_argument_parser.py
@@ -351,6 +351,57 @@ class EnumParser(ArgumentParser):
     return 'string enum'
 
 
+class EnumClassParser(ArgumentParser):
+  """Parser of an Enum class member."""
+
+  def __init__(self, enum_class):
+    """Initializes EnumParser.
+
+    Args:
+      enum_class: class, the Enum class with all possible flag values.
+
+    Raises:
+      TypeError: When enum_class is not a subclass of Enum.
+      ValueError: When enum_class is empty.
+    """
+    # Users must have an Enum class defined before using EnumClass flag.
+    # Therefore this dependency is guaranteed.
+    import enum
+
+    if not issubclass(enum_class, enum.Enum):
+      raise TypeError('{} is not a subclass of Enum.'.format(enum_class))
+    if not enum_class.__members__:
+      raise ValueError('enum_class cannot be empty, but "{}" is empty.'
+                       .format(enum_class))
+
+    super(EnumClassParser, self).__init__()
+    self.enum_class = enum_class
+
+  def parse(self, argument):
+    """Determines validity of argument and returns the correct element of enum.
+
+    Args:
+      argument: str or Enum class member, the supplied flag value.
+
+    Returns:
+      The first matching Enum class member in Enum class.
+
+    Raises:
+      ValueError: Raised when argument didn't match anything in enum.
+    """
+    if isinstance(argument, self.enum_class):
+      return argument
+    if argument not in self.enum_class.__members__:
+      raise ValueError('value should be one of <%s>' %
+                       '|'.join(self.enum_class.__members__.keys()))
+    else:
+      return self.enum_class[argument]
+
+  def flag_type(self):
+    """See base class."""
+    return 'enum class'
+
+
 class ListSerializer(ArgumentSerializer):
 
   def __init__(self, list_sep):