Add multi enum class flag to absl.flags.

This change enables the use of enum.Enum objects in multi_* style labels.

Previously the multi_enum flag in absl.flags only accepted strings. The new multi_enum_class flag accepts (lists of) Enum class members.

The text format of lists of enums when serializing into XML differs from the list format in other multi_* types of flags. This is because all other multi_* types rely on repr() to turn a list into a string (since all other multi_* types have primitive contents). With multi_enum_class, repr() prints out a list of the repr()'s of Enum objects, which are rather unsightly for public consumption (e.g., they contain the integer value of the enum flags). This change does some acrobatics to ensure XML representation of flag usage looks reasonable for lists, but those rendered lists look a little different from repr()'ed lists.

This change fixes the way multi_* flags are serialized when saved for future re-ingestion via a command line. It makes each flag value appear on a separate line, by modifying the list separator of all MultiFlag and subclasses to be '\n'.

PiperOrigin-RevId: 224405076

3 files changed, 174 insertions, 0 deletions

diff --git a/absl/flags/tests/_flag_test.py b/absl/flags/tests/_flag_test.py
index 522711c..084b37b 100644
--- a/absl/flags/tests/_flag_test.py
+++ b/absl/flags/tests/_flag_test.py
@@ -151,5 +151,50 @@ class EnumClassFlagTest(parameterized.TestCase):
       _flag.EnumClassFlag('fruit', 'BANANA', 'help', Fruit)
 
 
+class MultiEnumClassFlagTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('NoHelpSupplied', '', '<APPLE|ORANGE>: (no help available);\n    '
+       'repeat this option to specify a list of values'),
+      ('WithHelpSupplied', 'Type of fruit.',
+       '<APPLE|ORANGE>: Type of fruit.;\n    '
+       'repeat this option to specify a list of values'))
+  def test_help_text(self, helptext_input, helptext_output):
+    f = _flag.MultiEnumClassFlag('fruit', None, helptext_input, Fruit)
+    self.assertEqual(helptext_output, f.help)
+
+  def test_requires_enum(self):
+    with self.assertRaises(TypeError):
+      _flag.MultiEnumClassFlag('fruit', None, 'help', ['apple', 'orange'])
+
+  def test_requires_non_empty_enum_class(self):
+    with self.assertRaises(ValueError):
+      _flag.MultiEnumClassFlag('empty', None, 'help', EmptyEnum)
+
+  def test_accepts_literal_default(self):
+    f = _flag.MultiEnumClassFlag('fruit', Fruit.APPLE, 'A sample enum flag.',
+                                 Fruit)
+    self.assertListEqual([Fruit.APPLE], f.value)
+
+  def test_accepts_list_of_literal_default(self):
+    f = _flag.MultiEnumClassFlag('fruit', [Fruit.APPLE, Fruit.ORANGE],
+                                 'A sample enum flag.', Fruit)
+    self.assertListEqual([Fruit.APPLE, Fruit.ORANGE], f.value)
+
+  def test_accepts_string_default(self):
+    f = _flag.MultiEnumClassFlag('fruit', 'ORANGE', 'A sample enum flag.',
+                                 Fruit)
+    self.assertListEqual([Fruit.ORANGE], f.value)
+
+  def test_accepts_list_of_string_default(self):
+    f = _flag.MultiEnumClassFlag('fruit', ['ORANGE', 'APPLE'],
+                                 'A sample enum flag.', Fruit)
+    self.assertListEqual([Fruit.ORANGE, Fruit.APPLE], f.value)
+
+  def test_default_value_does_not_exist(self):
+    with self.assertRaises(_exceptions.IllegalFlagValueError):
+      _flag.MultiEnumClassFlag('fruit', 'BANANA', 'help', Fruit)
+
+
 if __name__ == '__main__':
   absltest.main()
diff --git a/absl/flags/tests/flags_helpxml_test.py b/absl/flags/tests/flags_helpxml_test.py
index 8324268..33b1e84 100644
--- a/absl/flags/tests/flags_helpxml_test.py
+++ b/absl/flags/tests/flags_helpxml_test.py
@@ -362,6 +362,49 @@ class FlagCreateXMLDOMElement(absltest.TestCase):
         '</flag>\n')
     self._check_flag_help_in_xml('flavours', 'tool', expected_output)
 
+  def test_flag_help_in_xml_multi_enum_class_singleton_default(self):
+    class Fruit(enum.Enum):
+      ORANGE = 0
+      BANANA = 1
+
+    flags.DEFINE_multi_enum_class('fruit', ['ORANGE'],
+                                  Fruit,
+                                  'The fruit flag.', flag_values=self.fv)
+    expected_output = (
+        '<flag>\n'
+        '  <file>tool</file>\n'
+        '  <name>fruit</name>\n'
+        '  <meaning>&lt;ORANGE|BANANA&gt;: The fruit flag.;\n'
+        '    repeat this option to specify a list of values</meaning>\n'
+        '  <default>ORANGE</default>\n'
+        '  <current>ORANGE</current>\n'
+        '  <type>multi enum class</type>\n'
+        '  <enum_value>ORANGE</enum_value>\n'
+        '  <enum_value>BANANA</enum_value>\n'
+        '</flag>\n')
+    self._check_flag_help_in_xml('fruit', 'tool', expected_output)
+
+  def test_flag_help_in_xml_multi_enum_class_list_default(self):
+    class Fruit(enum.Enum):
+      ORANGE = 0
+      BANANA = 1
+
+    flags.DEFINE_multi_enum_class('fruit', ['ORANGE', 'BANANA'],
+                                  Fruit,
+                                  'The fruit flag.', flag_values=self.fv)
+    expected_output = (
+        '<flag>\n'
+        '  <file>tool</file>\n'
+        '  <name>fruit</name>\n'
+        '  <meaning>&lt;ORANGE|BANANA&gt;: The fruit flag.;\n'
+        '    repeat this option to specify a list of values</meaning>\n'
+        '  <default>ORANGE,BANANA</default>\n'
+        '  <current>ORANGE,BANANA</current>\n'
+        '  <type>multi enum class</type>\n'
+        '  <enum_value>ORANGE</enum_value>\n'
+        '  <enum_value>BANANA</enum_value>\n'
+        '</flag>\n')
+    self._check_flag_help_in_xml('fruit', 'tool', expected_output)
 
 # The next EXPECTED_HELP_XML_* constants are parts of a template for
 # the expected XML output from WriteHelpInXMLFormatTest below.  When
diff --git a/absl/flags/tests/flags_test.py b/absl/flags/tests/flags_test.py
index 9c5d81f..90094d3 100644
--- a/absl/flags/tests/flags_test.py
+++ b/absl/flags/tests/flags_test.py
@@ -1309,6 +1309,92 @@ class MultiEnumFlagsTest(absltest.TestCase):
         FLAGS, argv)
 
 
+class MultiEnumClassFlagsTest(absltest.TestCase):
+
+  def test_define_results_in_registered_flag_with_none(self):
+    fv = flags.FlagValues()
+    enum_defaults = None
+    flags.DEFINE_multi_enum_class('fruit',
+                                  enum_defaults, Fruit,
+                                  'Enum option that can occur multiple times',
+                                  flag_values=fv)
+    fv.mark_as_parsed()
+
+    self.assertIsNone(fv.fruit)
+
+  def test_define_results_in_registered_flag_with_string(self):
+    fv = flags.FlagValues()
+    enum_defaults = 'apple'
+    flags.DEFINE_multi_enum_class('fruit',
+                                  enum_defaults, Fruit,
+                                  'Enum option that can occur multiple times',
+                                  flag_values=fv)
+    fv.mark_as_parsed()
+
+    self.assertListEqual(fv.fruit, [Fruit.apple])
+
+  def test_define_results_in_registered_flag_with_enum(self):
+    fv = flags.FlagValues()
+    enum_defaults = Fruit.APPLE
+    flags.DEFINE_multi_enum_class('fruit',
+                                  enum_defaults, Fruit,
+                                  'Enum option that can occur multiple times',
+                                  flag_values=fv)
+    fv.mark_as_parsed()
+
+    self.assertListEqual(fv.fruit, [Fruit.APPLE])
+
+  def test_define_results_in_registered_flag_with_string_list(self):
+    fv = flags.FlagValues()
+    enum_defaults = ['apple', 'APPLE']
+    flags.DEFINE_multi_enum_class('fruit',
+                                  enum_defaults, Fruit,
+                                  'Enum option that can occur multiple times',
+                                  flag_values=fv)
+    fv.mark_as_parsed()
+
+    self.assertListEqual(fv.fruit, [Fruit.apple, Fruit.APPLE])
+
+  def test_define_results_in_registered_flag_with_enum_list(self):
+    fv = flags.FlagValues()
+    enum_defaults = [Fruit.APPLE, Fruit.orange]
+    flags.DEFINE_multi_enum_class('fruit',
+                                  enum_defaults, Fruit,
+                                  'Enum option that can occur multiple times',
+                                  flag_values=fv)
+    fv.mark_as_parsed()
+
+    self.assertListEqual(fv.fruit, [Fruit.APPLE, Fruit.orange])
+
+  def test_from_command_line_returns_multiple(self):
+    fv = flags.FlagValues()
+    enum_defaults = [Fruit.APPLE]
+    flags.DEFINE_multi_enum_class('fruit',
+                                  enum_defaults, Fruit,
+                                  'Enum option that can occur multiple times',
+                                  flag_values=fv)
+    argv = ('./program', '--fruit=apple', '--fruit=orange')
+    fv(argv)
+    self.assertListEqual(fv.fruit, [Fruit.apple, Fruit.orange])
+
+  def test_bad_multi_enum_class_flags_from_definition(self):
+    with self.assertRaisesRegex(
+        flags.IllegalFlagValueError,
+        'flag --fruit=INVALID: value should be one of <apple|orange|APPLE>'):
+      flags.DEFINE_multi_enum_class('fruit', ['INVALID'], Fruit, 'desc')
+
+  def test_bad_multi_enum_class_flags_from_commandline(self):
+    fv = flags.FlagValues()
+    enum_defaults = [Fruit.APPLE]
+    flags.DEFINE_multi_enum_class('fruit', enum_defaults, Fruit, 'desc',
+                                  flag_values=fv)
+    argv = ('./program', '--fruit=INVALID')
+    with self.assertRaisesRegex(
+        flags.IllegalFlagValueError,
+        'flag --fruit=INVALID: value should be one of <apple|orange|APPLE>'):
+      fv(argv)
+
+
 class UnicodeFlagsTest(absltest.TestCase):
   """Testing proper unicode support for flags."""