diff --git a/bitmask/bitmask.py b/bitmask/bitmask.py index 32dcb35..e95a7da 100644 --- a/bitmask/bitmask.py +++ b/bitmask/bitmask.py @@ -3,7 +3,7 @@ """Utilities for manipulating bitmasks.""" -import bitmask.util as util +from bitmask.util import type_name, fullname class Bitmask: """Generic bitmask, which represents multiple Enum values. @@ -65,7 +65,7 @@ class Bitmask: @value.setter def value(self, value): if not issubclass(type(value), int): - raise TypeError(f"value must be an integer (got '{type(value)}')") + raise TypeError(f"value must be an integer") self._value = value def __contains__(self, item): @@ -79,7 +79,7 @@ class Bitmask: elif issubclass(type(item), self.AllFlags): return bool(self.value & item) else: - raise TypeError(f"item must be an {type(self)} or {self.AllFlags} (got '{type(item)}')") + raise TypeError(f"item must be {type_name(self)} or {type_name(self.AllFlags)}") def __iter__(self): """Return list of enabled flags.""" @@ -97,13 +97,13 @@ class Bitmask: return '|'.join([flag.name for flag in self]) or "0" def __repr__(self): - enum_name = util.fullname(self.AllFlags(0)) + enum_name = fullname(self.AllFlags(0)) args = ', '.join( [enum_name] + [f"{enum_name}.{flag.name}" for flag in self] ) - return f"{util.fullname(self)}({args})" + return f"{fullname(self)}({args})" def __eq__(self, other): """Check equality.""" @@ -123,7 +123,7 @@ class Bitmask: bitmask, given the initial value and the flag. """ if not issubclass(type(flag), self.AllFlags): - raise TypeError(f"can only add {self.AllFlags} (not '{type(flag)}') to {type(self)}") + raise TypeError(f"can only apply {type_name(self.AllFlags)} to {type_name(self)}") self.value = op(self.value, flag) def __mask_op(self, other, op): @@ -142,7 +142,7 @@ class Bitmask: elif issubclass(type(other), self.AllFlags): new_bitmask._flag_op(other, op) else: - raise TypeError(f"can only apply {type(self)} or {self.AllFlags} (not '{type(other)}') to {type(self)}") + raise TypeError(f"can only apply {type_name(self)} or {type_name(self.AllFlags)} to {type_name(self)}") return new_bitmask @@ -231,7 +231,7 @@ class Bitmask: TypeError: `flag` is not a single Enum value. """ if not issubclass(type(flag), self._AllFlags): - raise TypeError(f"can only discard {self.AllFlags} (not '{type(flag)}') from {type(self)}") + raise TypeError(f"can only discard {type_name(self._AllFlags)} from {type_name(self)}") return self._flag_op(flag, lambda a, b : a & ~b) diff --git a/bitmask/util.py b/bitmask/util.py index ec44b9a..6c4cf30 100644 --- a/bitmask/util.py +++ b/bitmask/util.py @@ -3,6 +3,8 @@ """Miscellaneous utilities for formatting exceptions, etc.""" +import enum + def fullname(obj): """Get the full class name of an object, including module. @@ -10,3 +12,20 @@ def fullname(obj): String with class name, which can be used in an `eval()`. """ return f"{obj.__class__.__qualname__}" + +def type_name(obj): + """Get the short type of an object. + + Returns: + String with human-readable type of the object. + + Example: + + import bitmask.util as util + util.type_name(1) + >>> int + """ + if issubclass(type(obj), enum.EnumMeta): + return obj.__name__ + else: + return obj.__class__.__name__ diff --git a/tests/test_bitmask.py b/tests/test_bitmask.py index 60eb61f..42539b7 100644 --- a/tests/test_bitmask.py +++ b/tests/test_bitmask.py @@ -24,6 +24,8 @@ def test_eq(): assert Bitmask(Desc, Desc.SMALL) != Desc.SMALL assert Bitmask(Desc, Desc.SMALL) != Desc.ROUND assert Bitmask(Desc) != Desc.ROUND + assert Bitmask(Desc) != "Hello World!" + assert Bitmask(Desc) != 0 def test_repr(): """Ensure evaluating __repr__ creates an identical object.""" @@ -45,6 +47,9 @@ def test_add(): mask.add(Desc.ROUND) assert mask == Bitmask(Desc, Desc.ROUND) + with pytest.raises(TypeError, match="can only apply Desc to Bitmask"): + mask.add(1) + def test_add_operator(): """Test + operator.""" # Individual flags @@ -154,7 +159,7 @@ def test_discard(): empty_mask = Bitmask(Desc) empty_mask.discard(Desc.SMALL) assert empty_mask == Bitmask(Desc) - with pytest.raises(TypeError): + with pytest.raises(TypeError, match="can only discard Desc from Bitmask"): empty_mask.discard(Bitmask(Desc, Desc.SMALL)) def test_subtract(): @@ -220,9 +225,9 @@ def test_value(): assert mask == Bitmask(Desc) mask.value = 1 assert mask == Bitmask(Desc, Desc.SMALL) - with pytest.raises(TypeError): + with pytest.raises(TypeError, match="value must be an integer"): mask.value = 1j - with pytest.raises(TypeError): + with pytest.raises(TypeError, match="value must be an integer"): mask.value = 2.5 def test_contains(): @@ -239,8 +244,9 @@ def test_contains(): # Bitmasks assert mask in mask assert Bitmask(Desc, Desc.SMALL, Desc.FUNKY) in mask + assert Bitmask(Desc, Desc.SMALL, Desc.FUNKY, Desc.ROUND) not in mask assert Bitmask(Desc, Desc.FUNKY) in mask - with pytest.raises(TypeError): + with pytest.raises(TypeError, match="item must be Bitmask or Desc"): x = 1 in mask def test_iter():