allow enum classes as mixins

ethanfurman · ethanfurman · commit fe119661a15f · 2018-09-21T01:34:49.000-07:00
diff --git a/Lib/enum.py b/Lib/enum.py
@@ -490,37 +490,13 @@ def _find_data_type(bases):
                             continue
                         return base
 
-        # check that zero or one concrete data type has been used, and only one
-        # Enum type has been used, and that that Enum has no members
-        member_type = first_enum = None
-        if len(bases) == 1:
-            member_type = _find_data_type(bases) or object
-            first_enum = bases[0]
-        elif len(bases) == 2:
-            mixin, first_enum = bases
-            if not issubclass(first_enum, Enum) or issubclass(mixin, Enum):
-                raise TypeError("new enumerations must be created as "
-                        "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
-            # search for a concrete data type in mixin or first_enum
-            member_type = _find_data_type(bases) or object
-        else:
-            # more than two bases
-            *mixins, first_enum = bases
-            if (not issubclass(first_enum, Enum)
-                    or any([issubclass(m, Enum) for m in mixins])
-                ):
-                raise TypeError("new enumerations must be created as "
-                        "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
-            *mixins, data_type = mixins
-            # verify that none of the mixins are a data type
-            if _find_data_type(mixins):
-                raise TypeError("new enumerations must be created as "
-                        "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
-            # check if data type is one, or base Enum has one
-            member_type = _find_data_type(bases[-2:]) or object
-
-        # finally, double check that we are not subclassing a class with existing
-        # enumeration members
+        # ensure final parent class is an Enum derivative, find any concrete
+        # data type, and check that Enum has no members
+        first_enum = bases[-1]
+        if not issubclass(first_enum, Enum):
+            raise TypeError("new enumerations should be created as "
+                    "`EnumName([mixin_type, ...] [data_type,] enum_type)`")
+        member_type = _find_data_type(bases) or object
         if first_enum._member_names_:
             raise TypeError("Cannot extend enumerations")
 
diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py
@@ -1756,6 +1756,14 @@ def MAX(cls):
         class StrMixin:
             def __str__(self):
                 return self._name_.lower()
+        class SomeEnum(Enum):
+            def behavior(self):
+                return 'booyah'
+        class AnotherEnum(Enum):
+            def behavior(self):
+                return 'nuhuh!'
+            def social(self):
+                return "what's up?"
         class Color(MaxMixin, Enum):
             RED = auto()
             GREEN = auto()
@@ -1783,6 +1791,56 @@ class Color(StrMixin, MaxMixin, Enum):
         self.assertEqual(Color.BLUE.value, 3)
         self.assertEqual(Color.MAX, 3)
         self.assertEqual(str(Color.BLUE), 'blue')
+        class CoolColor(StrMixin, SomeEnum, Enum):
+            RED = auto()
+            GREEN = auto()
+            BLUE = auto()
+        self.assertEqual(CoolColor.RED.value, 1)
+        self.assertEqual(CoolColor.GREEN.value, 2)
+        self.assertEqual(CoolColor.BLUE.value, 3)
+        self.assertEqual(str(CoolColor.BLUE), 'blue')
+        self.assertEqual(CoolColor.RED.behavior(), 'booyah')
+        class CoolerColor(StrMixin, AnotherEnum, Enum):
+            RED = auto()
+            GREEN = auto()
+            BLUE = auto()
+        self.assertEqual(CoolerColor.RED.value, 1)
+        self.assertEqual(CoolerColor.GREEN.value, 2)
+        self.assertEqual(CoolerColor.BLUE.value, 3)
+        self.assertEqual(str(CoolerColor.BLUE), 'blue')
+        self.assertEqual(CoolerColor.RED.behavior(), 'nuhuh!')
+        self.assertEqual(CoolerColor.RED.social(), "what's up?")
+        class CoolestColor(StrMixin, SomeEnum, AnotherEnum):
+            RED = auto()
+            GREEN = auto()
+            BLUE = auto()
+        self.assertEqual(CoolestColor.RED.value, 1)
+        self.assertEqual(CoolestColor.GREEN.value, 2)
+        self.assertEqual(CoolestColor.BLUE.value, 3)
+        self.assertEqual(str(CoolestColor.BLUE), 'blue')
+        self.assertEqual(CoolestColor.RED.behavior(), 'booyah')
+        self.assertEqual(CoolestColor.RED.social(), "what's up?")
+        class ConfusedColor(StrMixin, AnotherEnum, SomeEnum):
+            RED = auto()
+            GREEN = auto()
+            BLUE = auto()
+        self.assertEqual(ConfusedColor.RED.value, 1)
+        self.assertEqual(ConfusedColor.GREEN.value, 2)
+        self.assertEqual(ConfusedColor.BLUE.value, 3)
+        self.assertEqual(str(ConfusedColor.BLUE), 'blue')
+        self.assertEqual(ConfusedColor.RED.behavior(), 'nuhuh!')
+        self.assertEqual(ConfusedColor.RED.social(), "what's up?")
+        class ReformedColor(StrMixin, IntEnum, SomeEnum, AnotherEnum):
+            RED = auto()
+            GREEN = auto()
+            BLUE = auto()
+        self.assertEqual(ReformedColor.RED.value, 1)
+        self.assertEqual(ReformedColor.GREEN.value, 2)
+        self.assertEqual(ReformedColor.BLUE.value, 3)
+        self.assertEqual(str(ReformedColor.BLUE), 'blue')
+        self.assertEqual(ReformedColor.RED.behavior(), 'booyah')
+        self.assertEqual(ConfusedColor.RED.social(), "what's up?")
+        self.assertTrue(issubclass(ReformedColor, int))
 
 
 class TestOrder(unittest.TestCase):
