"""\ ========= Copyright ========= - Copyright: 2013 Ethan Furman -- All rights reserved. - Author: Ethan Furman - Version: 0.9 - Contact: ethan@stoneleaf.us Permissions: See the Python Contributor Licencse Agreement. """ from collections import OrderedDict import sys __all__ = ('Enum', 'IntEnum') class StealthProperty(): """ Returns the value of the instance, or the virtual attribute on the class. """ # Used to provide access to the `name` and `value` properties of enum # members while keeping some measure of protection from modification, # while still allowing for an enumeration to have members named `name` # and `value`. This works because enumeration members are not set # directely on the enum class -- __getattr__ is used to look them up. def __init__(self, fget=None, fset=None, fdel=None, doc=None): self.fget = fget self.fset = fset self.fdel = fdel self.__doc__ = doc or fget.__doc__ def __call__(self, func): self.fget = func self.name = func.name if not self.__doc__: self.__doc__ = self.fget.__doc__ def __get__(self, obj, objtype=None): if obj is None: return getattr(objtype, self.name) if self.fget is None: raise AttributeError("unreadable attribute") return self.fget(obj) def __set__(self, obj, value): if self.fset is None: raise AttributeError("can't set attribute") self.fset(obj, value) def __delete__(self, obj): if self.fdel is None: raise AttributeError("can't delete attribute") self.fdel(obj) def setter(self, func): self.fset = func return self def deleter(self, func): self.fdel = func return self class _EnumDict(dict): """Keeps track of definition order of the enum items. EnumMeta will use the names found in self._enum_names as the enumeration member names.""" def __init__(self): super().__init__() self._enum_names = [] def __setitem__(self, key, something): """Changes anything not dundered or that doesn't have __get__. All non-dunder names are also checked to ensure that enum items are not replaced with methods, and methods are not replaced with enum items. """ if key[:2] == key[-2:] == '__' or hasattr(something, '__get__'): if key in self._enum_names: # overwriting an enum with a method? self._enum_names.remove(key) else: if key in self._enum_names: raise TypeError('Attempted to reuse key: %s' % key) self._enum_names.append(key) dict.__setitem__(self, key, something) Enum = None # dummy value until replaced class EnumMeta(type): """Metaclass for Enum Pure enumerations can take any value, but the enumeration item will not compare equal to its value. Each enumeration item is created only once and is accessible as a virtual attribute of its class. Enumeration types can also be created by mixing in another data type; the resulting psuedenums must all be of that mixed-in type. """ @classmethod def __prepare__(metacls, cls, bases): return _EnumDict() def __new__(metacls, cls, bases, classdict): # an Enum class is final once enumeration items have been defined; # it cannot be mixed with other data types if it has an inhereted # __new__ unless a new __new__ is defined (or the resulting class # will fail). obj_type = first_enum = None # double check that we are not subclassing a class with existing # enumeration members; while we're at it, see if any other data # type has been mixed in so we can use the correct __new__ if bases: for base in bases: if (base is not Enum and issubclass(base, Enum) and base._enum_names): raise TypeError("Cannot extend enumerations") if not issubclass(base, Enum): raise TypeError("new enumerations must be created as " "`ClassName([mixin_type,] enum_type)`") # get correct mixin_type (either mixin type of Enum subclass, # or first base if last base is Enum if not issubclass(bases[0], Enum): obj_type = bases[0] # first data type first_enum = bases[-1] # enum type else: for base in bases[0].__mro__: # most common: (IntEnum, int, Enum, object) # possible: (, , # , , # ) if issubclass(base, Enum): if first_enum is None: first_enum = base else: if obj_type is None: obj_type = base else: obj_type = object first_enum = Enum # now find the correct __new__, checking to see of one was defined # by the user; also check earlier enum classes in case a __new__ was # saved as __new_member__ __new__ = classdict.get('__new__', None) # should __new__ be saved as __new_member__ later? save_new = __new__ is not None if __new__ is None: # check all possibles for __new_member__ before falling back to # __new__ for method in ('__new_member__', '__new__'): for possible in (obj_type, first_enum): target = getattr(possible, method, None) if target not in (None, None.__new__, object.__new__, Enum.__new__): __new__ = target break if __new__ is not None: break else: __new__ = object.__new__ # if a non-object.__new__ is used then whatever value/tuple was # assigned to the enum member name will be passed to __new__ and to the # new enum member's __init__ if __new__ is object.__new__: use_args = False else: use_args = True # save enum items into separate mapping so they don't get baked into # the new class name_value = {k: classdict[k] for k in classdict._enum_names} for name in classdict._enum_names: del classdict[name] enum_map = OrderedDict() enum_class = type.__new__(metacls, cls, bases, classdict) enum_names= [] enum_class._enum_names = enum_names # enum names in definition order enum_class._enum_map = enum_map # name:value map # instantiate them, checking for duplicates as we go # we instantiate first instead of checking for duplicates first in case # a custom __new__ is doing something funky with the values -- such as # auto-numbering ;) for e in classdict._enum_names: value = name_value[e] if not isinstance(value, tuple): args = (value, ) else: args = value if obj_type is tuple: # special case for tuple enums args = (args, ) # wrap it one more time if not use_args: enum_item = __new__(enum_class) else: enum_item = __new__(enum_class, *args) enum_item.__init__(*args) enum_item._name = e if not hasattr(enum_item, '_value'): enum_item._value = value # look for any duplicate values, and, if found, use the already # created enum item instead of the new one so `is` will work # (i.e. Color.green is Color.grene) for name, canonical_enum in enum_map.items(): if canonical_enum.value == enum_item._value: enum_item = canonical_enum break else: enum_names.append(e) enum_map[e] = enum_item # double check that repr and friends are not the mixin's or various # things break (such as pickle) for name in ('__repr__', '__str__', '__getnewargs__'): class_method = getattr(enum_class, name) obj_method = getattr(obj_type, name, None) enum_method = getattr(first_enum, name, None) if obj_method is not None and obj_method is class_method: setattr(enum_class, name, enum_method) # replace any other __new__ with our own (as long as Enum is not None, # anyway) -- again, this is to support pickle if Enum is not None: # if the user defined their own __new__, save it before it gets # clobbered in case they subclass later if save_new: enum_class.__new_member__ = __new__ enum_class.__new__ = Enum.__new__ return enum_class def __call__(cls, value, names=None, *, module_name=None, type=None): if names is None: # simple value lookup return cls.__new__(cls, value) # otherwise, we're creating a new Enum type class_name = value # better name for a name than value ;) metacls = cls.__class__ bases = (cls, ) if type is None else (type, cls) classdict = metacls.__prepare__(class_name, bases) if isinstance(names, str): names = names.replace(',', ' ').split() if isinstance(names, (tuple, list)) and isinstance(names[0], str): names = [(e, i) for (i, e) in enumerate(names, 1)] # otherwise names better be an iterable of (name, value) or a mapping for item in names: if isinstance(item, str): e, v = item, names[item] else: e, v = item classdict[e] = v enum_class = metacls.__new__(metacls, class_name, bases, classdict) # TODO: replace the frame hack if a blessed way to know the calling # module is ever developed if module_name is None: try: module_name = sys._getframe(1).f_globals['__name__'] except (AttributeError, ValueError): pass if module_name is not None: enum_class.__module__ = module_name return enum_class def __contains__(cls, enum_item): return isinstance(enum_item, cls) and enum_item.name in cls._enum_map def __dir__(self): return (['__class__', '__doc__', '__members__'] + list(self.__members__)) @property def __members__(cls): return cls._enum_map.copy() def __getattr__(cls, name): """Return the enum item matching `name`""" if name[:2] == name[-2:] == '__': raise AttributeError(name) try: return cls._enum_map[name] except KeyError: raise AttributeError(name) from None def __getitem__(cls, name): return cls._enum_map[name] def __iter__(cls): return (cls._enum_map[name] for name in cls._enum_names) def __len__(cls): return len(cls._enum_names) def __repr__(cls): return "" % cls.__name__ class Enum(metaclass=EnumMeta): """valueless, unordered enumeration class""" # no actual assignments are made as it is a chicken-and-egg problem # with the metaclass, which checks for the Enum class specifically def __new__(cls, value): # all enum instances are actually created during class construction # without calling this method; this method is called by the metaclass' # __call__ (i.e. Color(3) ), and by pickle if type(value) is cls: return value # by-value search for a matching enum member for enum_name in cls._enum_names: enum_member = cls._enum_map[enum_name] if enum_member.value == value: return enum_member raise ValueError("%s is not a valid %s" % (value, cls.__name__)) def __repr__(self): return "<%s.%s: %r>" % ( self.__class__.__name__, self._name, self._value) def __str__(self): return "%s.%s" % (self.__class__.__name__, self._name) def __dir__(self): return (['__class__', '__doc__', 'name', 'value']) def __eq__(self, other): if type(other) is self.__class__: return self is other return NotImplemented def __getnewargs__(self): return (self._value, ) def __hash__(self): return hash(self._name) @StealthProperty def name(self): return self._name @StealthProperty def value(self): return self._value class IntEnum(int, Enum): """Enum where members are also (and must be) ints""" if __name__ == '__main__': print("that's all, folks!")