Allow usage of EnumDict without setting a private attr. Add tests.

Author: encukou

Lib/enum.py

@@ -346,6 +346,7 @@ def __init__(self):
         self._last_values = []
         self._ignore = []
         self._auto_called = False
+        self._cls_name = None
 
     def __setitem__(self, key, value):
         """
@@ -356,7 +357,7 @@ def __setitem__(self, key, value):
 
         Single underscore (sunder) names are reserved.
         """
-        if _is_private(self._cls_name, key):
+        if self._cls_name is not None and _is_private(self._cls_name, key):
             # do nothing, name will be a normal attribute
             pass
         elif _is_sunder(key):
@@ -404,7 +405,7 @@ def __setitem__(self, key, value):
             value = value.value
         elif _is_descriptor(value):
             pass
-        elif _is_internal_class(self._cls_name, value):
+        elif self._cls_name is not None and _is_internal_class(self._cls_name, value):
             # do nothing, name will be a normal attribute
             pass
         else:

Lib/test/test_enum.py

@@ -14,7 +14,7 @@
 from enum import Enum, EnumMeta, IntEnum, StrEnum, EnumType, Flag, IntFlag, unique, auto
 from enum import STRICT, CONFORM, EJECT, KEEP, _simple_enum, _test_simple_enum
 from enum import verify, UNIQUE, CONTINUOUS, NAMED_FLAGS, ReprEnum
-from enum import member, nonmember, _iter_bits_lsb
+from enum import member, nonmember, _iter_bits_lsb, EnumDict
 from io import StringIO
 from pickle import dumps, loads, PicklingError, HIGHEST_PROTOCOL
 from test import support
@@ -5414,6 +5414,37 @@ def test_convert_repr_and_str(self):
         self.assertEqual(format(test_type.CONVERT_STRING_TEST_NAME_A), '5')
 
 
+class TestEnumDict(unittest.TestCase):
+    def test_enum_dict_in_metaclass(self):
+        """Test that EnumDict is usable as a class namespace"""
+        class Meta(type):
+            @classmethod
+            def __prepare__(metacls, cls, bases, **kwds):
+                return EnumDict()
+
+        class MyClass(metaclass=Meta):
+            a = 1
+
+            with self.assertRaises(TypeError):
+                a = 2  # duplicate
+
+            with self.assertRaises(ValueError):
+                _a_sunder_ = 3
+
+    def test_enum_dict_standalone(self):
+        """Test that EnumDict is usable on its own"""
+        enumdict = EnumDict()
+        enumdict['a'] = 1
+
+        with self.assertRaises(TypeError):
+            enumdict['a'] = 'other value'
+
+        # Only MutableMapping interface is overridden for now.
+        # If this starts passing, update the documentation.
+        enumdict |= {'a': 'other value'}
+        self.assertEqual(enumdict['a'], 'other value')
+
+
 # helpers
 
 def enum_dir(cls):