diff options
| -rw-r--r-- | Lib/enum.py | 57 | ||||
| -rw-r--r-- | Lib/test/test_enum.py | 16 |
2 files changed, 52 insertions, 21 deletions
diff --git a/Lib/enum.py b/Lib/enum.py index 98a8966f5eb..d10b9961598 100644 --- a/Lib/enum.py +++ b/Lib/enum.py @@ -279,9 +279,10 @@ class _proto_member: enum_member._sort_order_ = len(enum_class._member_names_) if Flag is not None and issubclass(enum_class, Flag): - enum_class._flag_mask_ |= value - if _is_single_bit(value): - enum_class._singles_mask_ |= value + if isinstance(value, int): + enum_class._flag_mask_ |= value + if _is_single_bit(value): + enum_class._singles_mask_ |= value enum_class._all_bits_ = 2 ** ((enum_class._flag_mask_).bit_length()) - 1 # If another member with the same value was already defined, the @@ -309,6 +310,7 @@ class _proto_member: elif ( Flag is not None and issubclass(enum_class, Flag) + and isinstance(value, int) and _is_single_bit(value) ): # no other instances found, record this member in _member_names_ @@ -1558,37 +1560,50 @@ class Flag(Enum, boundary=STRICT): def __bool__(self): return bool(self._value_) + def _get_value(self, flag): + if isinstance(flag, self.__class__): + return flag._value_ + elif self._member_type_ is not object and isinstance(flag, self._member_type_): + return flag + return NotImplemented + def __or__(self, other): - if isinstance(other, self.__class__): - other = other._value_ - elif self._member_type_ is not object and isinstance(other, self._member_type_): - other = other - else: + other_value = self._get_value(other) + if other_value is NotImplemented: return NotImplemented + + for flag in self, other: + if self._get_value(flag) is None: + raise TypeError(f"'{flag}' cannot be combined with other flags with |") value = self._value_ - return self.__class__(value | other) + return self.__class__(value | other_value) def __and__(self, other): - if isinstance(other, self.__class__): - other = other._value_ - elif self._member_type_ is not object and isinstance(other, self._member_type_): - other = other - else: + other_value = self._get_value(other) + if other_value is NotImplemented: return NotImplemented + + for flag in self, other: + if self._get_value(flag) is None: + raise TypeError(f"'{flag}' cannot be combined with other flags with &") value = self._value_ - return self.__class__(value & other) + return self.__class__(value & other_value) def __xor__(self, other): - if isinstance(other, self.__class__): - other = other._value_ - elif self._member_type_ is not object and isinstance(other, self._member_type_): - other = other - else: + other_value = self._get_value(other) + if other_value is NotImplemented: return NotImplemented + + for flag in self, other: + if self._get_value(flag) is None: + raise TypeError(f"'{flag}' cannot be combined with other flags with ^") value = self._value_ - return self.__class__(value ^ other) + return self.__class__(value ^ other_value) def __invert__(self): + if self._get_value(self) is None: + raise TypeError(f"'{self}' cannot be inverted") + if self._inverted_ is None: if self._boundary_ in (EJECT, KEEP): self._inverted_ = self.__class__(~self._value_) diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py index 61060f3dc29..cf3e042de1a 100644 --- a/Lib/test/test_enum.py +++ b/Lib/test/test_enum.py @@ -1048,6 +1048,22 @@ class TestPlainEnumFunction(_EnumTests, _PlainOutputTests, unittest.TestCase): class TestPlainFlagClass(_EnumTests, _PlainOutputTests, _FlagTests, unittest.TestCase): enum_type = Flag + def test_none_member(self): + class FlagWithNoneMember(Flag): + A = 1 + E = None + + self.assertEqual(FlagWithNoneMember.A.value, 1) + self.assertIs(FlagWithNoneMember.E.value, None) + with self.assertRaisesRegex(TypeError, r"'FlagWithNoneMember.E' cannot be combined with other flags with |"): + FlagWithNoneMember.A | FlagWithNoneMember.E + with self.assertRaisesRegex(TypeError, r"'FlagWithNoneMember.E' cannot be combined with other flags with &"): + FlagWithNoneMember.E & FlagWithNoneMember.A + with self.assertRaisesRegex(TypeError, r"'FlagWithNoneMember.E' cannot be combined with other flags with \^"): + FlagWithNoneMember.A ^ FlagWithNoneMember.E + with self.assertRaisesRegex(TypeError, r"'FlagWithNoneMember.E' cannot be inverted"): + ~FlagWithNoneMember.E + class TestPlainFlagFunction(_EnumTests, _PlainOutputTests, _FlagTests, unittest.TestCase): enum_type = Flag |