From 0063ff4e583505e69473caa978e476ea4c559b83 Mon Sep 17 00:00:00 2001 From: Ethan Furman Date: Mon, 21 Sep 2020 17:23:13 -0700 Subject: [PATCH] bpo-41816: add `StrEnum` (GH-22337) `StrEnum` ensures that its members were already strings, or intended to be strings. --- Doc/library/enum.rst | 38 +++++++++++++ Lib/enum.py | 32 ++++++++++- Lib/test/test_enum.py | 54 ++++++++++++------- .../2020-09-19-12-22-08.bpo-41816.ynynXJ.rst | 2 + 4 files changed, 104 insertions(+), 22 deletions(-) create mode 100644 Misc/NEWS.d/next/Library/2020-09-19-12-22-08.bpo-41816.ynynXJ.rst diff --git a/Doc/library/enum.rst b/Doc/library/enum.rst index 2f84be229bc..843d961afc4 100644 --- a/Doc/library/enum.rst +++ b/Doc/library/enum.rst @@ -44,6 +44,11 @@ helper, :class:`auto`. Base class for creating enumerated constants that are also subclasses of :class:`int`. +.. class:: StrEnum + + Base class for creating enumerated constants that are also + subclasses of :class:`str`. + .. class:: IntFlag Base class for creating enumerated constants that can be combined using @@ -601,6 +606,25 @@ However, they still can't be compared to standard :class:`Enum` enumerations:: [0, 1] +StrEnum +^^^^^^^ + +The second variation of :class:`Enum` that is provided is also a subclass of +:class:`str`. Members of a :class:`StrEnum` can be compared to strings; +by extension, string enumerations of different types can also be compared +to each other. :class:`StrEnum` exists to help avoid the problem of getting +an incorrect member:: + + >>> class Directions(StrEnum): + ... NORTH = 'north', # notice the trailing comma + ... SOUTH = 'south' + +Before :class:`StrEnum`, ``Directions.NORTH`` would have been the :class:`tuple` +``('north',)``. + +.. versionadded:: 3.10 + + IntFlag ^^^^^^^ @@ -1132,6 +1156,20 @@ all-uppercase names for members):: .. versionchanged:: 3.5 +Creating members that are mixed with other data types +""""""""""""""""""""""""""""""""""""""""""""""""""""" + +When subclassing other data types, such as :class:`int` or :class:`str`, with +an :class:`Enum`, all values after the `=` are passed to that data type's +constructor. For example:: + + >>> class MyEnum(IntEnum): + ... example = '11', 16 # '11' will be interpreted as a hexadecimal + ... # number + >>> MyEnum.example + + + Boolean value of ``Enum`` classes and members """"""""""""""""""""""""""""""""""""""""""""" diff --git a/Lib/enum.py b/Lib/enum.py index e8603a43420..589b17fd697 100644 --- a/Lib/enum.py +++ b/Lib/enum.py @@ -4,7 +4,7 @@ from types import MappingProxyType, DynamicClassAttribute __all__ = [ 'EnumMeta', - 'Enum', 'IntEnum', 'Flag', 'IntFlag', + 'Enum', 'IntEnum', 'StrEnum', 'Flag', 'IntFlag', 'auto', 'unique', ] @@ -688,7 +688,35 @@ class Enum(metaclass=EnumMeta): class IntEnum(int, Enum): - """Enum where members are also (and must be) ints""" + """ + Enum where members are also (and must be) ints + """ + + +class StrEnum(str, Enum): + """ + Enum where members are also (and must be) strings + """ + + def __new__(cls, *values): + if len(values) > 3: + raise TypeError('too many arguments for str(): %r' % (values, )) + if len(values) == 1: + # it must be a string + if not isinstance(values[0], str): + raise TypeError('%r is not a string' % (values[0], )) + if len(values) > 1: + # check that encoding argument is a string + if not isinstance(values[1], str): + raise TypeError('encoding must be a string, not %r' % (values[1], )) + if len(values) > 2: + # check that errors argument is a string + if not isinstance(values[2], str): + raise TypeError('errors must be a string, not %r' % (values[2], )) + value = str(*values) + member = str.__new__(cls, value) + member._value_ = value + return member def _reduce_ex_by_name(self, proto): diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py index 3f39073f5d5..8e84d929429 100644 --- a/Lib/test/test_enum.py +++ b/Lib/test/test_enum.py @@ -5,7 +5,7 @@ import sys import unittest import threading from collections import OrderedDict -from enum import Enum, IntEnum, EnumMeta, Flag, IntFlag, unique, auto +from enum import Enum, IntEnum, StrEnum, EnumMeta, Flag, IntFlag, unique, auto from io import StringIO from pickle import dumps, loads, PicklingError, HIGHEST_PROTOCOL from test import support @@ -48,14 +48,9 @@ except Exception as exc: FlagStooges = exc # for pickle test and subclass tests -try: - class StrEnum(str, Enum): - 'accepts only string values' - class Name(StrEnum): - BDFL = 'Guido van Rossum' - FLUFL = 'Barry Warsaw' -except Exception as exc: - Name = exc +class Name(StrEnum): + BDFL = 'Guido van Rossum' + FLUFL = 'Barry Warsaw' try: Question = Enum('Question', 'who what when where why', module=__name__) @@ -665,14 +660,13 @@ class TestEnum(unittest.TestCase): tau = 'Tau' self.assertTrue(phy.pi < phy.tau) - def test_strenum_inherited(self): - class StrEnum(str, Enum): - pass + def test_strenum_inherited_methods(self): class phy(StrEnum): pi = 'Pi' tau = 'Tau' self.assertTrue(phy.pi < phy.tau) - + self.assertEqual(phy.pi.upper(), 'PI') + self.assertEqual(phy.tau.count('a'), 1) def test_intenum(self): class WeekDay(IntEnum): @@ -2014,13 +2008,6 @@ class TestEnum(unittest.TestCase): self.assertTrue(issubclass(ReformedColor, int)) def test_multiple_inherited_mixin(self): - class StrEnum(str, Enum): - def __new__(cls, *args, **kwargs): - for a in args: - if not isinstance(a, str): - raise TypeError("Enumeration '%s' (%s) is not" - " a string" % (a, type(a).__name__)) - return str.__new__(cls, *args, **kwargs) @unique class Decision1(StrEnum): REVERT = "REVERT" @@ -2043,6 +2030,33 @@ class TestEnum(unittest.TestCase): local_ls = {} exec(code, global_ns, local_ls) + def test_strenum(self): + class GoodStrEnum(StrEnum): + one = '1' + two = '2' + three = b'3', 'ascii' + four = b'4', 'latin1', 'strict' + with self.assertRaisesRegex(TypeError, '1 is not a string'): + class FirstFailedStrEnum(StrEnum): + one = 1 + two = '2' + with self.assertRaisesRegex(TypeError, "2 is not a string"): + class SecondFailedStrEnum(StrEnum): + one = '1' + two = 2, + three = '3' + with self.assertRaisesRegex(TypeError, '2 is not a string'): + class ThirdFailedStrEnum(StrEnum): + one = '1' + two = 2 + with self.assertRaisesRegex(TypeError, 'encoding must be a string, not %r' % (sys.getdefaultencoding, )): + class ThirdFailedStrEnum(StrEnum): + one = '1' + two = b'2', sys.getdefaultencoding + with self.assertRaisesRegex(TypeError, 'errors must be a string, not 9'): + class ThirdFailedStrEnum(StrEnum): + one = '1' + two = b'2', 'ascii', 9 class TestOrder(unittest.TestCase): diff --git a/Misc/NEWS.d/next/Library/2020-09-19-12-22-08.bpo-41816.ynynXJ.rst b/Misc/NEWS.d/next/Library/2020-09-19-12-22-08.bpo-41816.ynynXJ.rst new file mode 100644 index 00000000000..605c346f37a --- /dev/null +++ b/Misc/NEWS.d/next/Library/2020-09-19-12-22-08.bpo-41816.ynynXJ.rst @@ -0,0 +1,2 @@ +StrEnum added: it ensures that all members are already strings or string +candidates