diff --git a/choicesenum/enums.py b/choicesenum/enums.py index f83c8e8..8639c5b 100644 --- a/choicesenum/enums.py +++ b/choicesenum/enums.py @@ -5,8 +5,24 @@ from enum import Enum, EnumMeta +def is_member_factory(enum_member): + "Return a property that checks if the current enum is the expected one" + @property + def is_member(self): + return self == enum_member + return is_member + + class ChoicesMetaClass(EnumMeta): + def __new__(metacls, cls, bases, classdict): + enum_class = EnumMeta.__new__(metacls, cls, bases, classdict) + for name, enum_value in enum_class._member_map_.items(): + prop_name = 'is_{}'.format(name.lower()) + setattr(enum_class, prop_name, is_member_factory(enum_value)) + + return enum_class + def __contains__(cls, member): if not isinstance(member, cls): try: @@ -25,13 +41,6 @@ def __new__(cls, value, display=None): obj._display_ = display return obj - def __getattr__(self, item): - is_attr = 'is_' - if item.startswith(is_attr) and item in self._get_dynamic_property_names(): - search = item[len(is_attr):] - return search == self._name_.lower() - raise AttributeError("'{}' object has no attribute '{}'".format(type(self).__name__, item)) - def __str__(self): return str(self.value) @@ -75,14 +84,6 @@ def __gt__(self, other): def __ge__(self, other): return self.value >= self._get_value(other) - def __dir__(self): - return sorted(set( - dir(type(self)) + - list(self.__dict__.keys()) + - ['display', 'get_choices', ] + - list(self._get_dynamic_property_names()) - )) - def __json__(self): """ If you want json serialization, you have at least two options: @@ -114,14 +115,6 @@ def description(self): """ return self.display - @classmethod - def _get_dynamic_property_names(cls): - """ - Args: - cls (Enum): Enum class. - """ - return ('is_{}'.format(x._name_.lower()) for x in cls) - @classmethod def choices(cls): """ diff --git a/tests/test_choicesenum.py b/tests/test_choicesenum.py index 954c69b..629b72c 100644 --- a/tests/test_choicesenum.py +++ b/tests/test_choicesenum.py @@ -68,11 +68,16 @@ def test_dynamic_is_attr_of_undefined_enums_should_fail(colors): assert colors.RED.is_banana -@pytest.mark.parametrize('attr', ['RED', 'BLUE', 'GREEN', 'is_red', 'is_blue', 'is_green']) +@pytest.mark.parametrize('attr', ['is_red', 'is_blue', 'is_green']) def test_dynamic_is_attr_should_be_in_dir(colors, attr): assert attr in dir(colors.RED) +@pytest.mark.parametrize('attr', ['RED', 'BLUE', 'GREEN']) +def test_enum_type_name_should_be_in_dir(colors, attr): + assert attr in dir(colors) + + @pytest.mark.skipif(sys.version_info < (3, 0), reason="requires python3") def test_in_format_python3(colors): assert '{}'.format(colors.RED) == "#f00"