bpo-38659: [Enum] add _simple_enum decorator (GH-25285)

add:

_simple_enum decorator to transform a normal class into an enum
_test_simple_enum function to compare
_old_convert_ to enable checking _convert_ generated enums
_simple_enum takes a normal class and converts it into an enum:

@simple_enum(Enum)
class Color:
    RED = 1
    GREEN = 2
    BLUE = 3

_old_convert_ works much like _convert_ does, using the original logic:

# in a test file
import socket, enum
CheckedAddressFamily = enum._old_convert_(
        enum.IntEnum, 'AddressFamily', 'socket',
        lambda C: C.isupper() and C.startswith('AF_'),
        source=_socket,
        )

test_simple_enum takes a traditional enum and a simple enum and
compares the two:

# in the REPL or the same module as Color
class CheckedColor(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

_test_simple_enum(CheckedColor, Color)

_test_simple_enum(CheckedAddressFamily, socket.AddressFamily)

Any important differences will raise a TypeError

Author: ethanfurman

Doc/library/enum.rst

@@ -621,4 +621,3 @@ Utilites and Decorators
       Traceback (most recent call last):
       ...
       ValueError: duplicate values found in <enum 'Mistake'>: FOUR -> THREE
-

Lib/ast.py

@@ -27,7 +27,7 @@
 import sys
 from _ast import *
 from contextlib import contextmanager, nullcontext
-from enum import IntEnum, auto
+from enum import IntEnum, auto, _simple_enum
 
 
 def parse(source, filename='<unknown>', mode='exec', *,
@@ -636,7 +636,8 @@ class Param(expr_context):
 # We unparse those infinities to INFSTR.
 _INFSTR = "1e" + repr(sys.float_info.max_10_exp + 1)
 
-class _Precedence(IntEnum):
+@_simple_enum(IntEnum)
+class _Precedence:
     """Precedence table that originated from python grammar."""
 
     TUPLE = auto()

Lib/enum.py

@@ -1,8 +1,10 @@
-from enum import IntEnum
+from enum import IntEnum, _simple_enum
 
 __all__ = ['HTTPStatus']
 
-class HTTPStatus(IntEnum):
+
+@_simple_enum(IntEnum)
+class HTTPStatus:
     """HTTP status codes and reason phrases
 
     Status codes from the following RFCs are all observed:

Lib/http/__init__.py

@@ -26,14 +26,15 @@
 import marshal
 import re
 
-from enum import Enum
+from enum import StrEnum, _simple_enum
 from functools import cmp_to_key
 from dataclasses import dataclass
 from typing import Dict
 
 __all__ = ["Stats", "SortKey", "FunctionProfile", "StatsProfile"]
 
-class SortKey(str, Enum):
+@_simple_enum(StrEnum)
+class SortKey:
     CALLS = 'calls', 'ncalls'
     CUMULATIVE = 'cumulative', 'cumtime'
     FILENAME = 'filename', 'module'

Lib/pstats.py

@@ -143,7 +143,8 @@
 __version__ = "2.2.1"
 
 @enum.global_enum
-class RegexFlag(enum.IntFlag, boundary=enum.KEEP):
+@enum._simple_enum(enum.IntFlag, boundary=enum.KEEP)
+class RegexFlag:
     ASCII = A = sre_compile.SRE_FLAG_ASCII # assume ascii "locale"
     IGNORECASE = I = sre_compile.SRE_FLAG_IGNORECASE # ignore case
     LOCALE = L = sre_compile.SRE_FLAG_LOCALE # assume current 8-bit locale

Lib/re.py

@@ -94,6 +94,7 @@
 import os
 from collections import namedtuple
 from enum import Enum as _Enum, IntEnum as _IntEnum, IntFlag as _IntFlag
+from enum import _simple_enum, _test_simple_enum
 
 import _ssl             # if we can't import it, let the error propagate
 
@@ -155,7 +156,8 @@
 _SSLv2_IF_EXISTS = getattr(_SSLMethod, 'PROTOCOL_SSLv2', None)
 
 
-class TLSVersion(_IntEnum):
+@_simple_enum(_IntEnum)
+class TLSVersion:
     MINIMUM_SUPPORTED = _ssl.PROTO_MINIMUM_SUPPORTED
     SSLv3 = _ssl.PROTO_SSLv3
     TLSv1 = _ssl.PROTO_TLSv1
@@ -165,7 +167,8 @@ class TLSVersion(_IntEnum):
     MAXIMUM_SUPPORTED = _ssl.PROTO_MAXIMUM_SUPPORTED
 
 
-class _TLSContentType(_IntEnum):
+@_simple_enum(_IntEnum)
+class _TLSContentType:
     """Content types (record layer)
 
     See RFC 8446, section B.1
@@ -179,7 +182,8 @@ class _TLSContentType(_IntEnum):
     INNER_CONTENT_TYPE = 0x101
 
 
-class _TLSAlertType(_IntEnum):
+@_simple_enum(_IntEnum)
+class _TLSAlertType:
     """Alert types for TLSContentType.ALERT messages
 
     See RFC 8466, section B.2
@@ -220,7 +224,8 @@ class _TLSAlertType(_IntEnum):
     NO_APPLICATION_PROTOCOL = 120
 
 
-class _TLSMessageType(_IntEnum):
+@_simple_enum(_IntEnum)
+class _TLSMessageType:
     """Message types (handshake protocol)
 
     See RFC 8446, section B.3

Lib/ssl.py

@@ -1,6 +1,7 @@
 import ast
 import builtins
 import dis
+import enum
 import os
 import sys
 import types
@@ -698,6 +699,35 @@ def test_constant_as_name(self):
             with self.assertRaisesRegex(ValueError, f"Name node can't be used with '{constant}' constant"):
                 compile(expr, "<test>", "eval")
 
+    def test_precedence_enum(self):
+        class _Precedence(enum.IntEnum):
+            """Precedence table that originated from python grammar."""
+            TUPLE = enum.auto()
+            YIELD = enum.auto()           # 'yield', 'yield from'
+            TEST = enum.auto()            # 'if'-'else', 'lambda'
+            OR = enum.auto()              # 'or'
+            AND = enum.auto()             # 'and'
+            NOT = enum.auto()             # 'not'
+            CMP = enum.auto()             # '<', '>', '==', '>=', '<=', '!=',
+                                          # 'in', 'not in', 'is', 'is not'
+            EXPR = enum.auto()
+            BOR = EXPR                    # '|'
+            BXOR = enum.auto()            # '^'
+            BAND = enum.auto()            # '&'
+            SHIFT = enum.auto()           # '<<', '>>'
+            ARITH = enum.auto()           # '+', '-'
+            TERM = enum.auto()            # '*', '@', '/', '%', '//'
+            FACTOR = enum.auto()          # unary '+', '-', '~'
+            POWER = enum.auto()           # '**'
+            AWAIT = enum.auto()           # 'await'
+            ATOM = enum.auto()
+            def next(self):
+                try:
+                    return self.__class__(self + 1)
+                except ValueError:
+                    return self
+        enum._test_simple_enum(_Precedence, ast._Precedence)
+
 
 class ASTHelpers_Test(unittest.TestCase):
     maxDiff = None

Lib/test/test_ast.py

@@ -8,7 +8,7 @@
 import threading
 from collections import OrderedDict
 from enum import Enum, IntEnum, StrEnum, EnumType, Flag, IntFlag, unique, auto
-from enum import STRICT, CONFORM, EJECT, KEEP
+from enum import STRICT, CONFORM, EJECT, KEEP, _simple_enum, _test_simple_enum
 from io import StringIO
 from pickle import dumps, loads, PicklingError, HIGHEST_PROTOCOL
 from test import support
@@ -2511,10 +2511,13 @@ class Bizarre(Flag, boundary=KEEP):
             d = 6
         #
         self.assertRaisesRegex(ValueError, 'invalid value: 7', Iron, 7)
+        #
         self.assertIs(Water(7), Water.ONE|Water.TWO)
         self.assertIs(Water(~9), Water.TWO)
+        #
         self.assertEqual(Space(7), 7)
         self.assertTrue(type(Space(7)) is int)
+        #
         self.assertEqual(list(Bizarre), [Bizarre.c])
         self.assertIs(Bizarre(3), Bizarre.b)
         self.assertIs(Bizarre(6), Bizarre.d)
@@ -3053,16 +3056,20 @@ class Space(IntFlag, boundary=EJECT):
             EIGHT = 8
         self.assertIs(Space._boundary_, EJECT)
         #
+        #
         class Bizarre(IntFlag, boundary=KEEP):
             b = 3
             c = 4
             d = 6
         #
         self.assertRaisesRegex(ValueError, 'invalid value: 5', Iron, 5)
+        #
         self.assertIs(Water(7), Water.ONE|Water.TWO)
         self.assertIs(Water(~9), Water.TWO)
+        #
         self.assertEqual(Space(7), 7)
         self.assertTrue(type(Space(7)) is int)
+        #
         self.assertEqual(list(Bizarre), [Bizarre.c])
         self.assertIs(Bizarre(3), Bizarre.b)
         self.assertIs(Bizarre(6), Bizarre.d)
@@ -3577,6 +3584,41 @@ def test_inspect_classify_class_attrs(self):
         if failed:
             self.fail("result does not equal expected, see print above")
 
+    def test_test_simple_enum(self):
+        @_simple_enum(Enum)
+        class SimpleColor:
+            RED = 1
+            GREEN = 2
+            BLUE = 3
+        class CheckedColor(Enum):
+            RED = 1
+            GREEN = 2
+            BLUE = 3
+        self.assertTrue(_test_simple_enum(CheckedColor, SimpleColor) is None)
+        SimpleColor.GREEN._value_ = 9
+        self.assertRaisesRegex(
+                TypeError, "enum mismatch",
+                _test_simple_enum, CheckedColor, SimpleColor,
+                )
+        class CheckedMissing(IntFlag, boundary=KEEP):
+            SIXTY_FOUR = 64
+            ONE_TWENTY_EIGHT = 128
+            TWENTY_FORTY_EIGHT = 2048
+            ALL = 2048 + 128 + 64 + 12
+        CM = CheckedMissing
+        self.assertEqual(list(CheckedMissing), [CM.SIXTY_FOUR, CM.ONE_TWENTY_EIGHT, CM.TWENTY_FORTY_EIGHT])
+        #
+        @_simple_enum(IntFlag, boundary=KEEP)
+        class Missing:
+            SIXTY_FOUR = 64
+            ONE_TWENTY_EIGHT = 128
+            TWENTY_FORTY_EIGHT = 2048
+            ALL = 2048 + 128 + 64 + 12
+        M = Missing
+        self.assertEqual(list(CheckedMissing), [M.SIXTY_FOUR, M.ONE_TWENTY_EIGHT, M.TWENTY_FORTY_EIGHT])
+        #
+        _test_simple_enum(CheckedMissing, Missing)
+        
 
 class MiscTestCase(unittest.TestCase):
     def test__all__(self):
@@ -3592,6 +3634,13 @@ def test__all__(self):
 CONVERT_TEST_NAME_E = 5
 CONVERT_TEST_NAME_F = 5
 
+CONVERT_STRING_TEST_NAME_D = 5
+CONVERT_STRING_TEST_NAME_C = 5
+CONVERT_STRING_TEST_NAME_B = 5
+CONVERT_STRING_TEST_NAME_A = 5  # This one should sort first.
+CONVERT_STRING_TEST_NAME_E = 5
+CONVERT_STRING_TEST_NAME_F = 5
+
 class TestIntEnumConvert(unittest.TestCase):
     def test_convert_value_lookup_priority(self):
         test_type = enum.IntEnum._convert_(
@@ -3639,14 +3688,16 @@ def test_convert_raise(self):
                 filter=lambda x: x.startswith('CONVERT_TEST_'))
 
     def test_convert_repr_and_str(self):
+        # reset global constants, as previous tests could have converted the 
+        # integer values to enums
         module = ('test.test_enum', '__main__')[__name__=='__main__']
         test_type = enum.IntEnum._convert_(
                 'UnittestConvert',
                 module,
-                filter=lambda x: x.startswith('CONVERT_TEST_'))
-        self.assertEqual(repr(test_type.CONVERT_TEST_NAME_A), '%s.CONVERT_TEST_NAME_A' % module)
-        self.assertEqual(str(test_type.CONVERT_TEST_NAME_A), 'CONVERT_TEST_NAME_A')
-        self.assertEqual(format(test_type.CONVERT_TEST_NAME_A), '5')
+                filter=lambda x: x.startswith('CONVERT_STRING_TEST_'))
+        self.assertEqual(repr(test_type.CONVERT_STRING_TEST_NAME_A), '%s.CONVERT_STRING_TEST_NAME_A' % module)
+        self.assertEqual(str(test_type.CONVERT_STRING_TEST_NAME_A), 'CONVERT_STRING_TEST_NAME_A')
+        self.assertEqual(format(test_type.CONVERT_STRING_TEST_NAME_A), '5')
 
 # global names for StrEnum._convert_ test
 CONVERT_STR_TEST_2 = 'goodbye'