Commit 94306525 authored by Rahul Jha's avatar Rahul Jha Committed by Ethan Furman

bpo-33217: Raise TypeError for non-Enum lookups in Enums (GH-6651)

* bpo-33217: Raise TypeError for non-Enum lookups in Enums
parent 51a4743d
......@@ -976,7 +976,7 @@ Enum Classes
The :class:`EnumMeta` metaclass is responsible for providing the
:meth:`__contains__`, :meth:`__dir__`, :meth:`__iter__` and other methods that
allow one to do things with an :class:`Enum` class that fail on a typical
class, such as `list(Color)` or `some_var in Color`. :class:`EnumMeta` is
class, such as `list(Color)` or `some_enum_var in Color`. :class:`EnumMeta` is
responsible for ensuring that various other methods on the final :class:`Enum`
class are correct (such as :meth:`__new__`, :meth:`__getnewargs__`,
:meth:`__str__` and :meth:`__repr__`).
......
......@@ -303,6 +303,10 @@ class EnumMeta(type):
return cls._create_(value, names, module=module, qualname=qualname, type=type, start=start)
def __contains__(cls, member):
if not isinstance(member, Enum):
raise TypeError(
"unsupported operand type(s) for 'in': '%s' and '%s'" % (
type(member).__qualname__, cls.__class__.__qualname__))
return isinstance(member, cls) and member._name_ in cls._member_map_
def __delattr__(cls, attr):
......@@ -705,7 +709,9 @@ class Flag(Enum):
def __contains__(self, other):
if not isinstance(other, self.__class__):
return NotImplemented
raise TypeError(
"unsupported operand type(s) for 'in': '%s' and '%s'" % (
type(other).__qualname__, self.__class__.__qualname__))
return other._value_ & self._value_ == other._value_
def __repr__(self):
......
......@@ -325,7 +325,10 @@ class TestEnum(unittest.TestCase):
def test_contains(self):
Season = self.Season
self.assertIn(Season.AUTUMN, Season)
self.assertNotIn(3, Season)
with self.assertRaises(TypeError):
3 in Season
with self.assertRaises(TypeError):
'AUTUMN' in Season
val = Season(3)
self.assertIn(val, Season)
......@@ -1752,6 +1755,13 @@ class TestFlag(unittest.TestCase):
AC = 3
CE = 1<<19
class Color(Flag):
BLACK = 0
RED = 1
GREEN = 2
BLUE = 4
PURPLE = RED|BLUE
def test_str(self):
Perm = self.Perm
self.assertEqual(str(Perm.R), 'Perm.R')
......@@ -1954,7 +1964,21 @@ class TestFlag(unittest.TestCase):
test_pickle_dump_load(self.assertIs, FlagStooges.CURLY|FlagStooges.MOE)
test_pickle_dump_load(self.assertIs, FlagStooges)
def test_containment(self):
def test_contains(self):
Open = self.Open
Color = self.Color
self.assertFalse(Color.BLACK in Open)
self.assertFalse(Open.RO in Color)
with self.assertRaises(TypeError):
'BLACK' in Color
with self.assertRaises(TypeError):
'RO' in Open
with self.assertRaises(TypeError):
1 in Color
with self.assertRaises(TypeError):
1 in Open
def test_member_contains(self):
Perm = self.Perm
R, W, X = Perm
RW = R | W
......@@ -2072,6 +2096,13 @@ class TestIntFlag(unittest.TestCase):
AC = 3
CE = 1<<19
class Color(IntFlag):
BLACK = 0
RED = 1
GREEN = 2
BLUE = 4
PURPLE = RED|BLUE
def test_type(self):
Perm = self.Perm
Open = self.Open
......@@ -2340,7 +2371,23 @@ class TestIntFlag(unittest.TestCase):
self.assertEqual(len(lst), len(Thing))
self.assertEqual(len(Thing), 0, Thing)
def test_containment(self):
def test_contains(self):
Open = self.Open
Color = self.Color
self.assertTrue(Color.GREEN in Color)
self.assertTrue(Open.RW in Open)
self.assertFalse(Color.GREEN in Open)
self.assertFalse(Open.RW in Color)
with self.assertRaises(TypeError):
'GREEN' in Color
with self.assertRaises(TypeError):
'RW' in Open
with self.assertRaises(TypeError):
2 in Color
with self.assertRaises(TypeError):
2 in Open
def test_member_contains(self):
Perm = self.Perm
R, W, X = Perm
RW = R | W
......@@ -2359,6 +2406,8 @@ class TestIntFlag(unittest.TestCase):
self.assertFalse(R in WX)
self.assertFalse(W in RX)
self.assertFalse(X in RW)
with self.assertRaises(TypeError):
self.assertFalse('test' in RW)
def test_bool(self):
Perm = self.Perm
......
Raise :exc:`TypeError` when looking up non-Enum objects in Enum classes and
Enum members.
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment