This does not directly answer the question, but I wanted to add a more robust version of @Ethan Furman's AutoEnum class which uses the auto
enum function.
The implementation below works with Pydantic and does fuzzy-matching of values to the corresponding enum type.
Usage:
In [2]: class Weekday(AutoEnum): ## Assume AutoEnum class has been defined.
...: Monday = auto()
...: Tuesday = auto()
...: Wednesday = auto()
...: Thursday = auto()
...: Friday = auto()
...: Saturday = auto()
...: Sunday = auto()
...:
In [3]: Weekday('MONDAY') ## Fuzzy matching: case-insensitive
Out[3]: Monday
In [4]: Weekday(' MO NDAY') ## Fuzzy matching: ignores extra spaces
Out[4]: Monday
In [5]: Weekday('_M_onDa y') ## Fuzzy matching: ignores underscores
Out[5]: Monday
In [6]: %timeit Weekday('_M_onDay') ## Fuzzy matching takes ~1 microsecond.
1.15 µs ± 10.9 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
In [7]: %timeit Weekday.from_str('_M_onDay') ## You can further speedup matching using from_str (this is because _missing_ is not called)
736 ns ± 8.89 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)
In [8]: list(Weekday) ## Get all the enums
Out[8]: [Monday, Tuesday, Wednesday, Thursday, Friday, Saturday, Sunday]
In [9]: Weekday.Monday.matches('Tuesday') ## Check if a string matches a particular enum value
Out[9]: False
In [10]: Weekday.matches_any('__TUESDAY__') ## Check if a string matches any enum
Out[10]: True
In [11]: Weekday.Tuesday is Weekday(' Tuesday') and Weekday.Tuesday == Weekday('_Tuesday_') ## `is` and `==` work as expected
Out[11]: True
In [12]: Weekday.Tuesday == 'Tuesday' ## Strings don't match enum values, because strings aren't enums!
Out[12]: False
In [13]: Weekday.convert_keys({ ## Convert matching dict keys to an enum. Similar: .convert_list, .convert_set
'monday': 'alice',
'tuesday': 'bob',
'not_wednesday': 'charles',
'THURSDAY ': 'denise',
})
Out[13]:
{Monday: 'alice',
Tuesday: 'bob',
'not_wednesday': 'charles',
Thursday: 'denise'}
The code for AutoEnum
can be found below.
If you want to change the fuzzy-matching logic, then override the classmethod _normalize
(e.g. returning the input unchanged in _normalize
, will perform exact matching).
from typing import *
from enum import Enum, auto
class AutoEnum(str, Enum):
"""
Utility class which can be subclassed to create enums using auto().
Also provides utility methods for common enum operations.
"""
@classmethod
def _missing_(cls, enum_value: Any):
## Ref: https://stackoverflow.com/a/60174274/4900327
## This is needed to allow Pydantic to perform case-insensitive conversion to AutoEnum.
return cls.from_str(enum_value=enum_value, raise_error=True)
def _generate_next_value_(name, start, count, last_values):
return name
@property
def str(self) -> str:
return self.__str__()
def __repr__(self):
return self.__str__()
def __str__(self):
return self.name
def __hash__(self):
return hash(self.__class__.__name__ + '.' + self.name)
def __eq__(self, other):
return self is other
def __ne__(self, other):
return self is not other
def matches(self, enum_value: str) -> bool:
return self is self.from_str(enum_value, raise_error=False)
@classmethod
def matches_any(cls, enum_value: str) -> bool:
return cls.from_str(enum_value, raise_error=False) is not None
@classmethod
def does_not_match_any(cls, enum_value: str) -> bool:
return not cls.matches_any(enum_value)
@classmethod
def _initialize_lookup(cls):
if '_value2member_map_normalized_' not in cls.__dict__: ## Caching values for fast retrieval.
cls._value2member_map_normalized_ = {}
for e in list(cls):
normalized_e_name: str = cls._normalize(e.value)
if normalized_e_name in cls._value2member_map_normalized_:
raise ValueError(
f'Cannot register enum "{e.value}"; '
f'another enum with the same normalized name "{normalized_e_name}" already exists.'
)
cls._value2member_map_normalized_[normalized_e_name] = e
@classmethod
def from_str(cls, enum_value: str, raise_error: bool = True) -> Optional:
"""
Performs a case-insensitive lookup of the enum value string among the members of the current AutoEnum subclass.
:param enum_value: enum value string
:param raise_error: whether to raise an error if the string is not found in the enum
:return: an enum value which matches the string
:raises: ValueError if raise_error is True and no enum value matches the string
"""
if isinstance(enum_value, cls):
return enum_value
if enum_value is None and raise_error is False:
return None
if not isinstance(enum_value, str) and raise_error is True:
raise ValueError(f'Input should be a string; found type {type(enum_value)}')
cls._initialize_lookup()
enum_obj: Optional[AutoEnum] = cls._value2member_map_normalized_.get(cls._normalize(enum_value))
if enum_obj is None and raise_error is True:
raise ValueError(f'Could not find enum with value {enum_value}; available values are: {list(cls)}.')
return enum_obj
@classmethod
def _normalize(cls, x: str) -> str:
## Found to be faster than .translate() and re.sub() on Python 3.10.6
return str(x).replace(' ', '').replace('-', '').replace('_', '').lower()
@classmethod
def convert_keys(cls, d: Dict) -> Dict:
"""
Converts string dict keys to the matching members of the current AutoEnum subclass.
Leaves non-string keys untouched.
:param d: dict to transform
:return: dict with matching string keys transformed to enum values
"""
out_dict = {}
for k, v in d.items():
if isinstance(k, str) and cls.from_str(k, raise_error=False) is not None:
out_dict[cls.from_str(k, raise_error=False)] = v
else:
out_dict[k] = v
return out_dict
@classmethod
def convert_keys_to_str(cls, d: Dict) -> Dict:
"""
Converts dict keys of the current AutoEnum subclass to the matching string key.
Leaves other keys untouched.
:param d: dict to transform
:return: dict with matching keys of the current AutoEnum transformed to strings.
"""
out_dict = {}
for k, v in d.items():
if isinstance(k, cls):
out_dict[str(k)] = v
else:
out_dict[k] = v
return out_dict
@classmethod
def convert_values(
cls,
d: Union[Dict, Set, List, Tuple],
raise_error: bool = False
) -> Union[Dict, Set, List, Tuple]:
"""
Converts string values to the matching members of the current AutoEnum subclass.
Leaves non-string values untouched.
:param d: dict, set, list or tuple to transform.
:param raise_error: raise an error if unsupported type.
:return: data structure with matching string values transformed to enum values.
"""
if isinstance(d, dict):
return cls.convert_dict_values(d)
if isinstance(d, list):
return cls.convert_list(d)
if isinstance(d, tuple):
return tuple(cls.convert_list(d))
if isinstance(d, set):
return cls.convert_set(d)
if raise_error:
raise ValueError(f'Unrecognized data structure of type {type(d)}')
return d
@classmethod
def convert_dict_values(cls, d: Dict) -> Dict:
"""
Converts string dict values to the matching members of the current AutoEnum subclass.
Leaves non-string values untouched.
:param d: dict to transform
:return: dict with matching string values transformed to enum values
"""
out_dict = {}
for k, v in d.items():
if isinstance(v, str) and cls.from_str(v, raise_error=False) is not None:
out_dict[k] = cls.from_str(v, raise_error=False)
else:
out_dict[k] = v
return out_dict
@classmethod
def convert_list(cls, l: List) -> List:
"""
Converts string list itmes to the matching members of the current AutoEnum subclass.
Leaves non-string items untouched.
:param l: list to transform
:return: list with matching string items transformed to enum values
"""
out_list = []
for item in l:
if isinstance(item, str) and cls.matches_any(item):
out_list.append(cls.from_str(item))
else:
out_list.append(item)
return out_list
@classmethod
def convert_set(cls, s: Set) -> Set:
"""
Converts string list itmes to the matching members of the current AutoEnum subclass.
Leaves non-string items untouched.
:param s: set to transform
:return: set with matching string items transformed to enum values
"""
out_set = set()
for item in s:
if isinstance(item, str) and cls.matches_any(item):
out_set.add(cls.from_str(item))
else:
out_set.add(item)
return out_set
@classmethod
def convert_values_to_str(cls, d: Dict) -> Dict:
"""
Converts dict values of the current AutoEnum subclass to the matching string value.
Leaves other values untouched.
:param d: dict to transform
:return: dict with matching values of the current AutoEnum transformed to strings.
"""
out_dict = {}
for k, v in d.items():
if isinstance(v, cls):
out_dict[k] = str(v)
else:
out_dict[k] = v
return out_dict