From 8c25ee325086ff7029aeeb839679f92b580298ba Mon Sep 17 00:00:00 2001 From: Edward Wang Date: Tue, 26 Jun 2018 21:01:19 -0700 Subject: [PATCH 1/2] bpo-33976: support nested classes in Enum --- Doc/library/enum.rst | 2 +- Lib/enum.py | 36 ++++++++++++++++-- Lib/test/test_enum.py | 38 +++++++++++++++++++ .../2018-06-28-11-04-10.bpo-33976.wahmah.rst | 1 + 4 files changed, 72 insertions(+), 5 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2018-06-28-11-04-10.bpo-33976.wahmah.rst diff --git a/Doc/library/enum.rst b/Doc/library/enum.rst index 0394681c088e75..79c03f448f7159 100644 --- 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 index 04d8ec1fa872f8..854d6154c58656 100644 --- a/Lib/enum.py +++ b/Lib/enum.py @@ -94,7 +94,10 @@ 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])) @@ -102,10 +105,21 @@ def __setitem__(self, key, value): 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 index 97559712b1dc2c..f2f8e0284dbc96 100644 --- 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 new file mode 100644 index 00000000000000..43f91b41a5e785 --- /dev/null +++ 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. From bbab8f9b3194d185870059853c430b2532dc443e Mon Sep 17 00:00:00 2001 From: Edward Wang Date: Thu, 28 Jun 2018 15:56:46 -0700 Subject: [PATCH 2/2] Catch the case where unrelated classes share a prefix --- Lib/enum.py | 2 +- Lib/test/test_enum.py | 5 +++++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/Lib/enum.py b/Lib/enum.py index 854d6154c58656..4946a1da37c9cc 100644 --- a/Lib/enum.py +++ b/Lib/enum.py @@ -156,7 +156,7 @@ def __new__(metacls, cls, bases, classdict): # accessible in _EnumDict.__setitem__(). for key, member in classdict.items(): if isinstance(member, type): - if member.__qualname__.startswith(enum_class_qualname): + if member.__qualname__.startswith(enum_class_qualname + "."): classdict.remove_member(key) # remove any keys listed in _ignore_ diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py index f2f8e0284dbc96..2df4689cd4bb50 100644 --- a/Lib/test/test_enum.py +++ b/Lib/test/test_enum.py @@ -419,6 +419,11 @@ class MyTypes2(Enum): b = Bar self.assertEqual(MyTypes2.a.value, Foo) self.assertEqual(MyTypes2.b.value, Bar) + class SpamEnumNotInner: + pass + class SpamEnum(Enum): + spam = SpamEnumNotInner + self.assertEqual(SpamEnum.spam.value, SpamEnumNotInner) def test_nested_classes_in_enum(self): """Support locally-defined nested classes."""