Skip to content

Commit

Permalink
bpo-41816: add StrEnum (pythonGH-22337)
Browse files Browse the repository at this point in the history
`StrEnum` ensures that its members were already strings, or intended to
be strings.
  • Loading branch information
ethanfurman committed Sep 22, 2020
1 parent 68526fe commit 0063ff4
Show file tree
Hide file tree
Showing 4 changed files with 104 additions and 22 deletions.
38 changes: 38 additions & 0 deletions Doc/library/enum.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
^^^^^^^

Expand Down Expand Up @@ -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
<MyEnum.example: 17>


Boolean value of ``Enum`` classes and members
"""""""""""""""""""""""""""""""""""""""""""""

Expand Down
32 changes: 30 additions & 2 deletions Lib/enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@

__all__ = [
'EnumMeta',
'Enum', 'IntEnum', 'Flag', 'IntFlag',
'Enum', 'IntEnum', 'StrEnum', 'Flag', 'IntFlag',
'auto', 'unique',
]

Expand Down Expand Up @@ -688,7 +688,35 @@ def value(self):


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):
Expand Down
54 changes: 34 additions & 20 deletions Lib/test/test_enum.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
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
Expand Down Expand Up @@ -48,14 +48,9 @@ class FlagStooges(Flag):
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__)
Expand Down Expand Up @@ -665,14 +660,13 @@ class phy(str, Enum):
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):
Expand Down Expand Up @@ -2014,13 +2008,6 @@ class ReformedColor(StrMixin, IntEnum, SomeEnum, AnotherEnum):
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"
Expand All @@ -2043,6 +2030,33 @@ def test_empty_globals(self):
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):

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
StrEnum added: it ensures that all members are already strings or string
candidates

0 comments on commit 0063ff4

Please sign in to comment.