Add AutoEnum: automatically provides next value if missing. Issue 26988.

1 files changed, 323 insertions, 1 deletions

diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py
index 564c0e9..4a732f9 100644
--- a/Lib/test/test_enum.py
+++ b/Lib/test/test_enum.py
@@ -3,7 +3,7 @@ import inspect
 import pydoc
 import unittest
 from collections import OrderedDict
-from enum import Enum, IntEnum, EnumMeta, unique
+from enum import EnumMeta, Enum, IntEnum, AutoEnum, unique
 from io import StringIO
 from pickle import dumps, loads, PicklingError, HIGHEST_PROTOCOL
 from test import support
@@ -1570,6 +1570,328 @@ class TestEnum(unittest.TestCase):
         self.assertEqual(LabelledList.unprocessed, 1)
         self.assertEqual(LabelledList(1), LabelledList.unprocessed)
 
+    def test_ignore_as_str(self):
+        from datetime import timedelta
+        class Period(Enum, ignore='Period i'):
+            """
+            different lengths of time
+            """
+            def __new__(cls, value, period):
+                obj = object.__new__(cls)
+                obj._value_ = value
+                obj.period = period
+                return obj
+            Period = vars()
+            for i in range(367):
+                Period['Day%d' % i] = timedelta(days=i), 'day'
+            for i in range(53):
+                Period['Week%d' % i] = timedelta(days=i*7), 'week'
+            for i in range(13):
+                Period['Month%d' % i] = i, 'month'
+            OneDay = Day1
+            OneWeek = Week1
+        self.assertEqual(Period.Day7.value, timedelta(days=7))
+        self.assertEqual(Period.Day7.period, 'day')
+
+    def test_ignore_as_list(self):
+        from datetime import timedelta
+        class Period(Enum, ignore=['Period', 'i']):
+            """
+            different lengths of time
+            """
+            def __new__(cls, value, period):
+                obj = object.__new__(cls)
+                obj._value_ = value
+                obj.period = period
+                return obj
+            Period = vars()
+            for i in range(367):
+                Period['Day%d' % i] = timedelta(days=i), 'day'
+            for i in range(53):
+                Period['Week%d' % i] = timedelta(days=i*7), 'week'
+            for i in range(13):
+                Period['Month%d' % i] = i, 'month'
+            OneDay = Day1
+            OneWeek = Week1
+        self.assertEqual(Period.Day7.value, timedelta(days=7))
+        self.assertEqual(Period.Day7.period, 'day')
+
+    def test_new_with_no_value_and_int_base_class(self):
+        class NoValue(int, Enum):
+            def __new__(cls, value):
+                obj = int.__new__(cls, value)
+                obj.index = len(cls.__members__)
+                return obj
+            this = 1
+            that = 2
+        self.assertEqual(list(NoValue), [NoValue.this, NoValue.that])
+        self.assertEqual(NoValue.this, 1)
+        self.assertEqual(NoValue.this.value, 1)
+        self.assertEqual(NoValue.this.index, 0)
+        self.assertEqual(NoValue.that, 2)
+        self.assertEqual(NoValue.that.value, 2)
+        self.assertEqual(NoValue.that.index, 1)
+
+    def test_new_with_no_value(self):
+        class NoValue(Enum):
+            def __new__(cls, value):
+                obj = object.__new__(cls)
+                obj.index = len(cls.__members__)
+                return obj
+            this = 1
+            that = 2
+        self.assertEqual(list(NoValue), [NoValue.this, NoValue.that])
+        self.assertEqual(NoValue.this.value, 1)
+        self.assertEqual(NoValue.this.index, 0)
+        self.assertEqual(NoValue.that.value, 2)
+        self.assertEqual(NoValue.that.index, 1)
+
+
+class TestAutoNumber(unittest.TestCase):
+
+    def test_autonumbering(self):
+        class Color(AutoEnum):
+            red
+            green
+            blue
+        self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
+        self.assertEqual(Color.red.value, 1)
+        self.assertEqual(Color.green.value, 2)
+        self.assertEqual(Color.blue.value, 3)
+
+    def test_autointnumbering(self):
+        class Color(int, AutoEnum):
+            red
+            green
+            blue
+        self.assertTrue(isinstance(Color.red, int))
+        self.assertEqual(Color.green, 2)
+        self.assertTrue(Color.blue > Color.red)
+
+    def test_autonumbering_with_start(self):
+        class Color(AutoEnum, start=7):
+            red
+            green
+            blue
+        self.assertEqual(list(Color), [Color.red, Color.green, Color.blue])
+        self.assertEqual(Color.red.value, 7)
+        self.assertEqual(Color.green.value, 8)
+        self.assertEqual(Color.blue.value, 9)
+
+    def test_autonumbering_with_start_and_skip(self):
+        class Color(AutoEnum, start=7):
+            red
+            green
+            blue = 11
+            brown
+        self.assertEqual(list(Color), [Color.red, Color.green, Color.blue, Color.brown])
+        self.assertEqual(Color.red.value, 7)
+        self.assertEqual(Color.green.value, 8)
+        self.assertEqual(Color.blue.value, 11)
+        self.assertEqual(Color.brown.value, 12)
+
+
+    def test_badly_overridden_ignore(self):
+        with self.assertRaisesRegex(TypeError, "'int' object is not callable"):
+            class Color(AutoEnum):
+                _ignore_ = ()
+                red
+                green
+                blue
+                @property
+                def whatever(self):
+                    pass
+        with self.assertRaisesRegex(TypeError, "'int' object is not callable"):
+            class Color(AutoEnum, ignore=None):
+                red
+                green
+                blue
+                @property
+                def whatever(self):
+                    pass
+        with self.assertRaisesRegex(TypeError, "'int' object is not callable"):
+            class Color(AutoEnum, ignore='classmethod staticmethod'):
+                red
+                green
+                blue
+                @property
+                def whatever(self):
+                    pass
+
+    def test_property(self):
+        class Color(AutoEnum):
+            red
+            green
+            blue
+            @property
+            def cap_name(self):
+                return self.name.title()
+        self.assertEqual(Color.blue.cap_name, 'Blue')
+
+    def test_magic_turns_off(self):
+        with self.assertRaisesRegex(NameError, "brown"):
+            class Color(AutoEnum):
+                red
+                green
+                blue
+                @property
+                def cap_name(self):
+                    return self.name.title()
+                brown
+
+        with self.assertRaisesRegex(NameError, "rose"):
+            class Color(AutoEnum):
+                red
+                green
+                blue
+                def hello(self):
+                    print('Hello!  My serial is %s.' % self.value)
+                rose
+
+        with self.assertRaisesRegex(NameError, "cyan"):
+            class Color(AutoEnum):
+                red
+                green
+                blue
+                def __init__(self, *args):
+                    pass
+                cyan
+
+
+class TestGenerateMethod(unittest.TestCase):
+
+    def test_autonaming(self):
+        class Color(Enum):
+            def _generate_next_value_(name, start, count, last_value):
+                return name
+            Red
+            Green
+            Blue
+        self.assertEqual(list(Color), [Color.Red, Color.Green, Color.Blue])
+        self.assertEqual(Color.Red.value, 'Red')
+        self.assertEqual(Color.Green.value, 'Green')
+        self.assertEqual(Color.Blue.value, 'Blue')
+
+    def test_autonamestr(self):
+        class Color(str, Enum):
+            def _generate_next_value_(name, start, count, last_value):
+                return name
+            Red
+            Green
+            Blue
+        self.assertTrue(isinstance(Color.Red, str))
+        self.assertEqual(Color.Green, 'Green')
+        self.assertTrue(Color.Blue < Color.Red)
+
+    def test_generate_as_staticmethod(self):
+        class Color(str, Enum):
+            @staticmethod
+            def _generate_next_value_(name, start, count, last_value):
+                return name.lower()
+            Red
+            Green
+            Blue
+        self.assertTrue(isinstance(Color.Red, str))
+        self.assertEqual(Color.Green, 'green')
+        self.assertTrue(Color.Blue < Color.Red)
+
+
+    def test_overridden_ignore(self):
+        with self.assertRaisesRegex(TypeError, "'str' object is not callable"):
+            class Color(Enum):
+                def _generate_next_value_(name, start, count, last_value):
+                    return name
+                _ignore_ = ()
+                red
+                green
+                blue
+                @property
+                def whatever(self):
+                    pass
+        with self.assertRaisesRegex(TypeError, "'str' object is not callable"):
+            class Color(Enum, ignore=None):
+                def _generate_next_value_(name, start, count, last_value):
+                    return name
+                red
+                green
+                blue
+                @property
+                def whatever(self):
+                    pass
+
+    def test_property(self):
+        class Color(Enum):
+            def _generate_next_value_(name, start, count, last_value):
+                return name
+            red
+            green
+            blue
+            @property
+            def upper_name(self):
+                return self.name.upper()
+        self.assertEqual(Color.blue.upper_name, 'BLUE')
+
+    def test_magic_turns_off(self):
+        with self.assertRaisesRegex(NameError, "brown"):
+            class Color(Enum):
+                def _generate_next_value_(name, start, count, last_value):
+                    return name
+                red
+                green
+                blue
+                @property
+                def cap_name(self):
+                    return self.name.title()
+                brown
+
+        with self.assertRaisesRegex(NameError, "rose"):
+            class Color(Enum):
+                def _generate_next_value_(name, start, count, last_value):
+                    return name
+                red
+                green
+                blue
+                def hello(self):
+                    print('Hello!  My value %s.' % self.value)
+                rose
+
+        with self.assertRaisesRegex(NameError, "cyan"):
+            class Color(Enum):
+                def _generate_next_value_(name, start, count, last_value):
+                    return name
+                red
+                green
+                blue
+                def __init__(self, *args):
+                    pass
+                cyan
+
+    def test_powers_of_two(self):
+        class Bits(Enum):
+            def _generate_next_value_(name, start, count, last_value):
+                return 2 ** count
+            one
+            two
+            four
+            eight
+        self.assertEqual(Bits.one.value, 1)
+        self.assertEqual(Bits.two.value, 2)
+        self.assertEqual(Bits.four.value, 4)
+        self.assertEqual(Bits.eight.value, 8)
+
+    def test_powers_of_two_as_int(self):
+        class Bits(int, Enum):
+            def _generate_next_value_(name, start, count, last_value):
+                return 2 ** count
+            one
+            two
+            four
+            eight
+        self.assertEqual(Bits.one, 1)
+        self.assertEqual(Bits.two, 2)
+        self.assertEqual(Bits.four, 4)
+        self.assertEqual(Bits.eight, 8)
+
 
 class TestUnique(unittest.TestCase):