bpo-33976: support nested classes in Enum

edwardcwang · edwardcwang · commit 8c25ee325086 · 2018-06-28T15:30:14.000-07:00
diff --git a/Doc/library/enum.rst b/Doc/library/enum.rst
@@ -207,7 +207,7 @@ return A::
 .. note::
 
     Attempting to create a member with the same name as an already
-    defined attribute (another member, a method, etc.) or attempting to create
+    defined attribute (another member, a method, a class, etc.) or attempting to create
     an attribute with the same name as a member is not allowed.
 
 
diff --git a/Lib/enum.py b/Lib/enum.py
@@ -94,18 +94,32 @@ def __setitem__(self, key, value):
             raise TypeError('Attempted to reuse key: %r' % key)
         elif key in self._ignore:
             pass
-        elif not _is_descriptor(value):
+        elif _is_descriptor(value):
+            # Don't treat methods, etc as enum values.
+            pass
+        else:
             if key in self:
                 # enum overwriting a descriptor?
                 raise TypeError('%r already defined as: %r' % (key, self[key]))
             if isinstance(value, auto):
                 if value.value == _auto_null:
                     value.value = self._generate_next_value(key, 1, len(self._member_names), self._last_values[:])
                 value = value.value
-            self._member_names.append(key)
-            self._last_values.append(value)
+            self.add_member(key, value)
         super().__setitem__(key, value)
 
+    def add_member(self, key, value):
+        """Add a member by key and value."""
+        self._member_names.append(key)
+        self._last_values.append(value)
+
+    def remove_member(self, key):
+        """Remove a member (reverses add_member() above) by key, if present."""
+        if key in self._member_names:
+            index = self._member_names.index(key)
+            del self._member_names[index]
+            del self._last_values[index]
+
 
 # Dummy value for Enum as EnumMeta explicitly checks for it, but of course
 # until EnumMeta finishes running the first time the Enum class doesn't exist.
@@ -130,7 +144,21 @@ def __new__(metacls, cls, bases, classdict):
         # cannot be mixed with other types (int, float, etc.) if it has an
         # inherited __new__ unless a new __new__ is defined (or the resulting
         # class will fail).
-        #
+
+        # Get __qualname__ for the class being created.
+        enum_class_qualname = super().__new__(metacls, cls, bases, classdict).__qualname__
+
+        # We want to avoid treating locally-defined nested classes as enum
+        # values, so we use __qualname__ to determine this.
+        # e.g. if a class Bar is defined inside an enum Foo, then say if
+        # enum_class.__qualname__ is Foo, then member.__qualname__ will be Bar.
+        # We have to do it here since __qualname__ of the new Enum isn't
+        # accessible in _EnumDict.__setitem__().
+        for key, member in classdict.items():
+            if isinstance(member, type):
+                if member.__qualname__.startswith(enum_class_qualname):
+                    classdict.remove_member(key)
+
         # remove any keys listed in _ignore_
         classdict.setdefault('_ignore_', []).append('_ignore_')
         ignore = classdict['_ignore_']
diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py
@@ -401,6 +401,44 @@ def red(self):
                 green = 2
                 blue = 3
 
+    def test_enum_of_types(self):
+        """Support using Enum to refer to types deliberately."""
+        class MyTypes(Enum):
+            i = int
+            f = float
+            s = str
+        self.assertEqual(MyTypes.i.value, int)
+        self.assertEqual(MyTypes.f.value, float)
+        self.assertEqual(MyTypes.s.value, str)
+        class Foo:
+            pass
+        class Bar:
+            pass
+        class MyTypes2(Enum):
+            a = Foo
+            b = Bar
+        self.assertEqual(MyTypes2.a.value, Foo)
+        self.assertEqual(MyTypes2.b.value, Bar)
+
+    def test_nested_classes_in_enum(self):
+        """Support locally-defined nested classes."""
+        class Outer(Enum):
+            a = 1
+            b = 2
+            class Inner(Enum):
+                foo = 10
+                bar = 11
+        self.assertTrue(isinstance(Outer.Inner, type))
+        self.assertEqual(Outer.a.value, 1)
+        self.assertEqual(Outer.Inner.foo.value, 10)
+        self.assertEqual(
+            list(Outer.Inner),
+            [Outer.Inner.foo, Outer.Inner.bar],
+            )
+        self.assertEqual(
+            list(Outer),
+            [Outer.a, Outer.b],
+            )
 
     def test_enum_with_value_name(self):
         class Huh(Enum):
diff --git a/Misc/NEWS.d/next/Library/2018-06-28-11-04-10.bpo-33976.wahmah.rst b/Misc/NEWS.d/next/Library/2018-06-28-11-04-10.bpo-33976.wahmah.rst
@@ -0,0 +1 @@
+Support nested classes in Enum. Patch by Edward Wang.

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	`+Support nested classes in Enum. Patch by Edward Wang.`