diff options
author | Ethan Furman <ethan@stoneleaf.us> | 2013-06-14 16:55:46 -0700 |
---|---|---|
committer | Ethan Furman <ethan@stoneleaf.us> | 2013-06-14 16:55:46 -0700 |
commit | 6fdb575288c9a12d3f1eeef9033d58753aa8ef38 (patch) | |
tree | 358a7d4aa57d0c789a85e8b6f9e5554cf21da4bd | |
parent | 44c11ed8833ebcff2c34478e715d67f1dbe62d0d (diff) | |
download | cpython-6fdb575288c9a12d3f1eeef9033d58753aa8ef38.tar.gz |
Closes issue 17947. Adds PEP-0435 (Adding an Enum type to the Python standard library).
Missing files added.
News entry added.
-rw-r--r-- | Doc/library/enum.rst | 542 | ||||
-rw-r--r-- | Lib/enum.py | 465 | ||||
-rw-r--r-- | Lib/test/test_enum.py | 921 | ||||
-rw-r--r-- | Misc/NEWS | 2 |
4 files changed, 1930 insertions, 0 deletions
diff --git a/Doc/library/enum.rst b/Doc/library/enum.rst new file mode 100644 index 0000000000..92c6bfbba1 --- /dev/null +++ b/Doc/library/enum.rst @@ -0,0 +1,542 @@ +:mod:`enum` --- Support for enumerations +======================================== + +.. module:: enum +.. :synopsis: enumerations are sets of symbolic names bound to unique, constant + values. +.. :moduleauthor:: Ethan Furman <ethan@stoneleaf.us> +.. :sectionauthor:: Barry Warsaw <barry@python.org>, +.. :sectionauthor:: Eli Bendersky <eliben@gmail.com>, +.. :sectionauthor:: Ethan Furman <ethan@stoneleaf.us> + +**Source code:** :source:`Lib/enum.py` + +---------------- + +An enumeration is a set of symbolic names (members) bound to unique, constant +values. Within an enumeration, the members can be compared by identity, and +the enumeration itself can be iterated over. + +This module defines two enumeration classes that can be used to define unique +sets of names and values: :class:`Enum` and :class:`IntEnum`. + +Creating an Enum +---------------- + +Enumerations are created using the :keyword:`class` syntax, which makes them +easy to read and write. An alternative creation method is described in +`Functional API`_. To define an enumeration, subclass :class:`Enum` as +follows:: + + >>> from enum import Enum + >>> class Color(Enum): + ... red = 1 + ... green = 2 + ... blue = 3 + +**A note on nomenclature**: we call :class:`Color` an *enumeration* (or *enum*) +and :attr:`Color.red`, :attr:`Color.green` are *enumeration members* (or +*enum members*). Enumeration members also have *values* (the value of +:attr:`Color.red` is ``1``, etc.) + +Enumeration members have human readable string representations:: + + >>> print(Color.red) + Color.red + +...while their ``repr`` has more information:: + + >>> print(repr(Color.red)) + <Color.red: 1> + +The *type* of an enumeration member is the enumeration it belongs to:: + + >>> type(Color.red) + <enum 'Color'> + >>> isinstance(Color.green, Color) + True + >>> + +Enum members also have a property that contains just their item name:: + + >>> print(Color.red.name) + red + +Enumerations support iteration, in definition order:: + + >>> class Shake(Enum): + ... vanilla = 7 + ... chocolate = 4 + ... cookies = 9 + ... mint = 3 + ... + >>> for shake in Shake: + ... print(shake) + ... + Shake.vanilla + Shake.chocolate + Shake.cookies + Shake.mint + +Enumeration members are hashable, so they can be used in dictionaries and sets:: + + >>> apples = {} + >>> apples[Color.red] = 'red delicious' + >>> apples[Color.green] = 'granny smith' + >>> apples == {Color.red: 'red delicious', Color.green: 'granny smith'} + True + + +Programmatic access to enumeration members +------------------------------------------ + +Sometimes it's useful to access members in enumerations programmatically (i.e. +situations where ``Color.red`` won't do because the exact color is not known +at program-writing time). ``Enum`` allows such access:: + + >>> Color(1) + <Color.red: 1> + >>> Color(3) + <Color.blue: 3> + +If you want to access enum members by *name*, use item access:: + + >>> Color['red'] + <Color.red: 1> + >>> Color['green'] + <Color.green: 2> + + +Duplicating enum members and values +----------------------------------- + +Having two enum members with the same name is invalid:: + + >>> class Shape(Enum): + ... square = 2 + ... square = 3 + ... + Traceback (most recent call last): + ... + TypeError: Attempted to reuse key: 'square' + +However, two enum members are allowed to have the same value. Given two members +A and B with the same value (and A defined first), B is an alias to A. By-value +lookup of the value of A and B will return A. By-name lookup of B will also +return A:: + + >>> class Shape(Enum): + ... square = 2 + ... diamond = 1 + ... circle = 3 + ... alias_for_square = 2 + ... + >>> Shape.square + <Shape.square: 2> + >>> Shape.alias_for_square + <Shape.square: 2> + >>> Shape(2) + <Shape.square: 2> + +Iterating over the members of an enum does not provide the aliases:: + + >>> list(Shape) + [<Shape.square: 2>, <Shape.diamond: 1>, <Shape.circle: 3>] + +The special attribute ``__members__`` is an ordered dictionary mapping names +to members. It includes all names defined in the enumeration, including the +aliases:: + + >>> for name, member in Shape.__members__.items(): + ... name, member + ... + ('square', <Shape.square: 2>) + ('diamond', <Shape.diamond: 1>) + ('circle', <Shape.circle: 3>) + ('alias_for_square', <Shape.square: 2>) + +The ``__members__`` attribute can be used for detailed programmatic access to +the enumeration members. For example, finding all the aliases:: + + >>> [name for name, member in Shape.__members__.items() if member.name != name] + ['alias_for_square'] + +Comparisons +----------- + +Enumeration members are compared by identity:: + + >>> Color.red is Color.red + True + >>> Color.red is Color.blue + False + >>> Color.red is not Color.blue + True + +Ordered comparisons between enumeration values are *not* supported. Enum +members are not integers (but see `IntEnum`_ below):: + + >>> Color.red < Color.blue + Traceback (most recent call last): + File "<stdin>", line 1, in <module> + TypeError: unorderable types: Color() < Color() + +Equality comparisons are defined though:: + + >>> Color.blue == Color.red + False + >>> Color.blue != Color.red + True + >>> Color.blue == Color.blue + True + +Comparisons against non-enumeration values will always compare not equal +(again, class:`IntEnum` was explicitly designed to behave differently, see +below):: + + >>> Color.blue == 2 + False + + +Allowed members and attributes of enumerations +---------------------------------------------- + +The examples above use integers for enumeration values. Using integers is +short and handy (and provided by default by the `Functional API`_), but not +strictly enforced. In the vast majority of use-cases, one doesn't care what +the actual value of an enumeration is. But if the value *is* important, +enumerations can have arbitrary values. + +Enumerations are Python classes, and can have methods and special methods as +usual. If we have this enumeration:: + + >>> class Mood(Enum): + ... funky = 1 + ... happy = 3 + ... + ... def describe(self): + ... # self is the member here + ... return self.name, self.value + ... + ... def __str__(self): + ... return 'my custom str! {0}'.format(self.value) + ... + ... @classmethod + ... def favorite_mood(cls): + ... # cls here is the enumeration + ... return cls.happy + +Then:: + + >>> Mood.favorite_mood() + <Mood.happy: 3> + >>> Mood.happy.describe() + ('happy', 3) + >>> str(Mood.funky) + 'my custom str! 1' + +The rules for what is allowed are as follows: _sunder_ names (starting and +ending with a single underscore) are reserved by enum and cannot be used; +all other attributes defined within an enumeration will become members of this +enumeration, with the exception of *__dunder__* names and descriptors (methods +are also descriptors). + +Note: if your enumeration defines :meth:`__new__` and/or :meth:`__init__` then +whatever value(s) were given to the enum member will be passed into those +methods. See `Planet`_ for an example. + + +Restricted subclassing of enumerations +-------------------------------------- + +Subclassing an enumeration is allowed only if the enumeration does not define +any members. So this is forbidden:: + + >>> class MoreColor(Color): + ... pink = 17 + Traceback (most recent call last): + ... + TypeError: Cannot extend enumerations + +But this is allowed:: + + >>> class Foo(Enum): + ... def some_behavior(self): + ... pass + ... + >>> class Bar(Foo): + ... happy = 1 + ... sad = 2 + ... + +Allowing subclassing of enums that define members would lead to a violation of +some important invariants of types and instances. On the other hand, it makes +sense to allow sharing some common behavior between a group of enumerations. +(See `OrderedEnum`_ for an example.) + + +Pickling +-------- + +Enumerations can be pickled and unpickled:: + + >>> from test.test_enum import Fruit + >>> from pickle import dumps, loads + >>> Fruit.tomato is loads(dumps(Fruit.tomato)) + True + +The usual restrictions for pickling apply: picklable enums must be defined in +the top level of a module, since unpickling requires them to be importable +from that module. + +.. warning:: + + In order to support the singleton nature of enumeration members, pickle + protocol version 2 or higher must be used. + + +Functional API +-------------- + +The :class:`Enum` class is callable, providing the following functional API:: + + >>> Animal = Enum('Animal', 'ant bee cat dog') + >>> Animal + <enum 'Animal'> + >>> Animal.ant + <Animal.ant: 1> + >>> Animal.ant.value + 1 + >>> list(Animal) + [<Animal.ant: 1>, <Animal.bee: 2>, <Animal.cat: 3>, <Animal.dog: 4>] + +The semantics of this API resemble :class:`namedtuple`. The first argument +of the call to :class:`Enum` is the name of the enumeration. + +The second argument is the *source* of enumeration member names. It can be a +whitespace-separated string of names, a sequence of names, a sequence of +2-tuples with key/value pairs, or a mapping (e.g. dictionary) of names to +values. The last two options enable assigning arbitrary values to +enumerations; the others auto-assign increasing integers starting with 1. A +new class derived from :class:`Enum` is returned. In other words, the above +assignment to :class:`Animal` is equivalent to:: + + >>> class Animals(Enum): + ... ant = 1 + ... bee = 2 + ... cat = 3 + ... dog = 4 + +Pickling enums created with the functional API can be tricky as frame stack +implementation details are used to try and figure out which module the +enumeration is being created in (e.g. it will fail if you use a utility +function in separate module, and also may not work on IronPython or Jython). +The solution is to specify the module name explicitly as follows:: + + >>> Animals = Enum('Animals', 'ant bee cat dog', module=__name__) + +Derived Enumerations +==================== + +IntEnum +------- + +A variation of :class:`Enum` is provided which is also a subclass of +:class:`int`. Members of an :class:`IntEnum` can be compared to integers; +by extension, integer enumerations of different types can also be compared +to each other:: + + >>> from enum import IntEnum + >>> class Shape(IntEnum): + ... circle = 1 + ... square = 2 + ... + >>> class Request(IntEnum): + ... post = 1 + ... get = 2 + ... + >>> Shape == 1 + False + >>> Shape.circle == 1 + True + >>> Shape.circle == Request.post + True + +However, they still can't be compared to standard :class:`Enum` enumerations:: + + >>> class Shape(IntEnum): + ... circle = 1 + ... square = 2 + ... + >>> class Color(Enum): + ... red = 1 + ... green = 2 + ... + >>> Shape.circle == Color.red + False + +:class:`IntEnum` values behave like integers in other ways you'd expect:: + + >>> int(Shape.circle) + 1 + >>> ['a', 'b', 'c'][Shape.circle] + 'b' + >>> [i for i in range(Shape.square)] + [0, 1] + +For the vast majority of code, :class:`Enum` is strongly recommended, +since :class:`IntEnum` breaks some semantic promises of an enumeration (by +being comparable to integers, and thus by transitivity to other +unrelated enumerations). It should be used only in special cases where +there's no other choice; for example, when integer constants are +replaced with enumerations and backwards compatibility is required with code +that still expects integers. + + +Others +------ + +While :class:`IntEnum` is part of the :mod:`enum` module, it would be very +simple to implement independently:: + + class IntEnum(int, Enum): + pass + +This demonstrates how similar derived enumerations can be defined; for example +a :class:`StrEnum` that mixes in :class:`str` instead of :class:`int`. + +Some rules: + +1. When subclassing :class:`Enum`, mix-in types must appear before + :class:`Enum` itself in the sequence of bases, as in the :class:`IntEnum` + example above. +2. While :class:`Enum` can have members of any type, once you mix in an + additional type, all the members must have values of that type, e.g. + :class:`int` above. This restriction does not apply to mix-ins which only + add methods and don't specify another data type such as :class:`int` or + :class:`str`. +3. When another data type is mixed in, the :attr:`value` attribute is *not the + same* as the enum member itself, although it is equivalant and will compare + equal. + + +Interesting examples +==================== + +While :class:`Enum` and :class:`IntEnum` are expected to cover the majority of +use-cases, they cannot cover them all. Here are recipes for some different +types of enumerations that can be used directly, or as examples for creating +one's own. + + +AutoNumber +---------- + +Avoids having to specify the value for each enumeration member:: + + >>> class AutoNumber(Enum): + ... def __new__(cls): + ... value = len(cls.__members__) + 1 + ... obj = object.__new__(cls) + ... obj._value = value + ... return obj + ... + >>> class Color(AutoNumber): + ... red = () + ... green = () + ... blue = () + ... + >>> Color.green.value == 2 + True + + +UniqueEnum +---------- + +Raises an error if a duplicate member name is found instead of creating an +alias:: + + >>> class UniqueEnum(Enum): + ... def __init__(self, *args): + ... cls = self.__class__ + ... if any(self.value == e.value for e in cls): + ... a = self.name + ... e = cls(self.value).name + ... raise ValueError( + ... "aliases not allowed in UniqueEnum: %r --> %r" + ... % (a, e)) + ... + >>> class Color(UniqueEnum): + ... red = 1 + ... green = 2 + ... blue = 3 + ... grene = 2 + Traceback (most recent call last): + ... + ValueError: aliases not allowed in UniqueEnum: 'grene' --> 'green' + + +OrderedEnum +----------- + +An ordered enumeration that is not based on :class:`IntEnum` and so maintains +the normal :class:`Enum` invariants (such as not being comparable to other +enumerations):: + + >>> class OrderedEnum(Enum): + ... def __ge__(self, other): + ... if self.__class__ is other.__class__: + ... return self._value >= other._value + ... return NotImplemented + ... def __gt__(self, other): + ... if self.__class__ is other.__class__: + ... return self._value > other._value + ... return NotImplemented + ... def __le__(self, other): + ... if self.__class__ is other.__class__: + ... return self._value <= other._value + ... return NotImplemented + ... def __lt__(self, other): + ... if self.__class__ is other.__class__: + ... return self._value < other._value + ... return NotImplemented + ... + >>> class Grade(OrderedEnum): + ... A = 5 + ... B = 4 + ... C = 3 + ... D = 2 + ... F = 1 + ... + >>> Grade.C < Grade.A + True + + +Planet +------ + +If :meth:`__new__` or :meth:`__init__` is defined the value of the enum member +will be passed to those methods:: + + >>> class Planet(Enum): + ... MERCURY = (3.303e+23, 2.4397e6) + ... VENUS = (4.869e+24, 6.0518e6) + ... EARTH = (5.976e+24, 6.37814e6) + ... MARS = (6.421e+23, 3.3972e6) + ... JUPITER = (1.9e+27, 7.1492e7) + ... SATURN = (5.688e+26, 6.0268e7) + ... URANUS = (8.686e+25, 2.5559e7) + ... NEPTUNE = (1.024e+26, 2.4746e7) + ... def __init__(self, mass, radius): + ... self.mass = mass # in kilograms + ... self.radius = radius # in meters + ... @property + ... def surface_gravity(self): + ... # universal gravitational constant (m3 kg-1 s-2) + ... G = 6.67300E-11 + ... return G * self.mass / (self.radius * self.radius) + ... + >>> Planet.EARTH.value + (5.976e+24, 6378140.0) + >>> Planet.EARTH.surface_gravity + 9.802652743337129 diff --git a/Lib/enum.py b/Lib/enum.py new file mode 100644 index 0000000000..775489bf95 --- /dev/null +++ b/Lib/enum.py @@ -0,0 +1,465 @@ +"""Python Enumerations""" + +import sys +from collections import OrderedDict +from types import MappingProxyType + +__all__ = ['Enum', 'IntEnum'] + + +class _RouteClassAttributeToGetattr: + """Route attribute access on a class to __getattr__. + + This is a descriptor, used to define attributes that act differently when + accessed through an instance and through a class. Instance access remains + normal, but access to an attribute through a class will be routed to the + class's __getattr__ method; this is done by raising AttributeError. + + """ + def __init__(self, fget=None): + self.fget = fget + + def __get__(self, instance, ownerclass=None): + if instance is None: + raise AttributeError() + return self.fget(instance) + + def __set__(self, instance, value): + raise AttributeError("can't set attribute") + + def __delete__(self, instance): + raise AttributeError("can't delete attribute") + + +def _is_dunder(name): + """Returns True if a __dunder__ name, False otherwise.""" + return (name[:2] == name[-2:] == '__' and + name[2:3] != '_' and + name[-3:-2] != '_') + + +def _is_sunder(name): + """Returns True if a _sunder_ name, False otherwise.""" + return (name[0] == name[-1] == '_' and + name[1:2] != '_' and + name[-2:-1] != '_') + + +def _make_class_unpicklable(cls): + """Make the given class un-picklable.""" + def _break_on_call_reduce(self): + raise TypeError('%r cannot be pickled' % self) + cls.__reduce__ = _break_on_call_reduce + cls.__module__ = '<unknown>' + + +class _EnumDict(dict): + """Keeps track of definition order of the enum items. + + EnumMeta will use the names found in self._member_names as the + enumeration member names. + + """ + def __init__(self): + super().__init__() + self._member_names = [] + + def __setitem__(self, key, value): + """Changes anything not dundered or that doesn't have __get__. + + If a descriptor is added with the same name as an enum member, the name + is removed from _member_names (this may leave a hole in the numerical + sequence of values). + + If an enum member name is used twice, an error is raised; duplicate + values are not checked for. + + Single underscore (sunder) names are reserved. + + """ + if _is_sunder(key): + raise ValueError('_names_ are reserved for future Enum use') + elif _is_dunder(key) or hasattr(value, '__get__'): + if key in self._member_names: + # overwriting an enum with a method? then remove the name from + # _member_names or it will become an enum anyway when the class + # is created + self._member_names.remove(key) + else: + if key in self._member_names: + raise TypeError('Attempted to reuse key: %r' % key) + self._member_names.append(key) + super().__setitem__(key, value) + + +# Dummy value for Enum as EnumMeta explicity checks for it, but of course until +# EnumMeta finishes running the first time the Enum class doesn't exist. This +# is also why there are checks in EnumMeta like `if Enum is not None` +Enum = None + + +class EnumMeta(type): + """Metaclass for Enum""" + @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 types (int, float, etc.) if it has an + # inherited __new__ unless a new __new__ is defined (or the resulting + # class will fail). + member_type, first_enum = metacls._get_mixins_(bases) + __new__, save_new, use_args = metacls._find_new_(classdict, member_type, + first_enum) + + # save enum items into separate mapping so they don't get baked into + # the new class + members = {k: classdict[k] for k in classdict._member_names} + for name in classdict._member_names: + del classdict[name] + + # check for illegal enum names (any others?) + invalid_names = set(members) & {'mro', } + if invalid_names: + raise ValueError('Invalid enum member name: {0}'.format( + ','.join(invalid_names))) + + # create our new Enum type + enum_class = super().__new__(metacls, cls, bases, classdict) + enum_class._member_names = [] # names in definition order + enum_class._member_map = OrderedDict() # name->value map + + # Reverse value->name map for hashable values. + enum_class._value2member_map = {} + + # check for a __getnewargs__, and if not present sabotage + # pickling, since it won't work anyway + if (member_type is not object and + member_type.__dict__.get('__getnewargs__') is None + ): + _make_class_unpicklable(enum_class) + + # 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 member_name in classdict._member_names: + value = members[member_name] + if not isinstance(value, tuple): + args = (value, ) + else: + args = value + if member_type is tuple: # special case for tuple enums + args = (args, ) # wrap it one more time + if not use_args: + enum_member = __new__(enum_class) + enum_member._value = value + else: + enum_member = __new__(enum_class, *args) + if not hasattr(enum_member, '_value'): + enum_member._value = member_type(*args) + enum_member._member_type = member_type + enum_member._name = member_name + enum_member.__init__(*args) + # If another member with the same value was already defined, the + # new member becomes an alias to the existing one. + for name, canonical_member in enum_class._member_map.items(): + if canonical_member.value == enum_member._value: + enum_member = canonical_member + break + else: + # Aliases don't appear in member names (only in __members__). + enum_class._member_names.append(member_name) + enum_class._member_map[member_name] = enum_member + try: + # This may fail if value is not hashable. We can't add the value + # to the map, and by-value lookups for this value will be + # linear. + enum_class._value2member_map[value] = enum_member + except TypeError: + pass + + # 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(member_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=None, type=None): + """Either returns an existing member, or creates a new enum class. + + This method is used both when an enum class is given a value to match + to an enumeration member (i.e. Color(3)) and for the functional API + (i.e. Color = Enum('Color', names='red green blue')). + + When used for the functional API: `module`, if set, will be stored in + the new class' __module__ attribute; `type`, if set, will be mixed in + as the first base class. + + Note: if `module` is not set this routine will attempt to discover the + calling module by walking the frame stack; if this is unsuccessful + the resulting class will not be pickleable. + + """ + if names is None: # simple value lookup + return cls.__new__(cls, value) + # otherwise, functional API: we're creating a new Enum type + return cls._create_(value, names, module=module, type=type) + + def __contains__(cls, member): + return isinstance(member, cls) and member.name in cls._member_map + + def __dir__(self): + return ['__class__', '__doc__', '__members__'] + self._member_names + + @property + def __members__(cls): + """Returns a mapping of member name->value. + + This mapping lists all enum members, including aliases. Note that this + is a read-only view of the internal mapping. + + """ + return MappingProxyType(cls._member_map) + + def __getattr__(cls, name): + """Return the enum member matching `name` + + We use __getattr__ instead of descriptors or inserting into the enum + class' __dict__ in order to support `name` and `value` being both + properties for enum members (which live in the class' __dict__) and + enum members themselves. + + """ + if _is_dunder(name): + raise AttributeError(name) + try: + return cls._member_map[name] + except KeyError: + raise AttributeError(name) from None + + def __getitem__(cls, name): + return cls._member_map[name] + + def __iter__(cls): + return (cls._member_map[name] for name in cls._member_names) + + def __len__(cls): + return len(cls._member_names) + + def __repr__(cls): + return "<enum %r>" % cls.__name__ + + def _create_(cls, class_name, names=None, *, module=None, type=None): + """Convenience method to create a new Enum class. + + `names` can be: + + * A string containing member names, separated either with spaces or + commas. Values are auto-numbered from 1. + * An iterable of member names. Values are auto-numbered from 1. + * An iterable of (member name, value) pairs. + * A mapping of member name -> value. + + """ + metacls = cls.__class__ + bases = (cls, ) if type is None else (type, cls) + classdict = metacls.__prepare__(class_name, bases) + + # special processing needed for names? + 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)] + + # Here, names is either an iterable of (name, value) or a mapping. + for item in names: + if isinstance(item, str): + member_name, member_value = item, names[item] + else: + member_name, member_value = item + classdict[member_name] = member_value + 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 is None: + try: + module = sys._getframe(2).f_globals['__name__'] + except (AttributeError, ValueError) as exc: + pass + if module is None: + _make_class_unpicklable(enum_class) + else: + enum_class.__module__ = module + + return enum_class + + @staticmethod + def _get_mixins_(bases): + """Returns the type for creating enum members, and the first inherited + enum class. + + bases: the tuple of bases that was given to __new__ + + """ + if not bases: + return object, Enum + + # 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__ + member_type = first_enum = None + for base in bases: + if (base is not Enum and + issubclass(base, Enum) and + base._member_names): + raise TypeError("Cannot extend enumerations") + # base is now the last base in bases + if not issubclass(base, Enum): + raise TypeError("new enumerations must be created as " + "`ClassName([mixin_type,] enum_type)`") + + # get correct mix-in type (either mix-in type of Enum subclass, or + # first base if last base is Enum) + if not issubclass(bases[0], Enum): + member_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: (<Enum 'AutoIntEnum'>, <Enum 'IntEnum'>, + # <class 'int'>, <Enum 'Enum'>, + # <class 'object'>) + if issubclass(base, Enum): + if first_enum is None: + first_enum = base + else: + if member_type is None: + member_type = base + + return member_type, first_enum + + @staticmethod + def _find_new_(classdict, member_type, first_enum): + """Returns the __new__ to be used for creating the enum members. + + classdict: the class dictionary given to __new__ + member_type: the data type whose __new__ will be used by default + first_enum: enumeration to check for an overriding __new__ + + """ + # 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 (member_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 + + return __new__, save_new, use_args + + +class Enum(metaclass=EnumMeta): + """Generic enumeration. + + Derive from this class to define new enumerations. + + """ + 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: + # For lookups like Color(Color.red) + return value + # by-value search for a matching enum member + # see if it's in the reverse mapping (for hashable values) + if value in cls._value2member_map: + return cls._value2member_map[value] + # not there, now do long search -- O(n) behavior + for member in cls._member_map.values(): + if member.value == value: + return 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) + + # _RouteClassAttributeToGetattr is 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 directly on the enum class -- __getattr__ is + # used to look them up. + + @_RouteClassAttributeToGetattr + def name(self): + return self._name + + @_RouteClassAttributeToGetattr + def value(self): + return self._value + + +class IntEnum(int, Enum): + """Enum where members are also (and must be) ints""" diff --git a/Lib/test/test_enum.py b/Lib/test/test_enum.py new file mode 100644 index 0000000000..75b2656862 --- /dev/null +++ b/Lib/test/test_enum.py @@ -0,0 +1,921 @@ +import enum +import unittest +from collections import OrderedDict +from pickle import dumps, loads, PicklingError +from enum import Enum, IntEnum + +# for pickle tests +try: + class Stooges(Enum): + LARRY = 1 + CURLY = 2 + MOE = 3 +except Exception as exc: + Stooges = exc + +try: + class IntStooges(int, Enum): + LARRY = 1 + CURLY = 2 + MOE = 3 +except Exception as exc: + IntStooges = exc + +try: + class FloatStooges(float, Enum): + LARRY = 1.39 + CURLY = 2.72 + MOE = 3.142596 +except Exception as exc: + FloatStooges = exc + +# for pickle test and subclass tests +try: + class StrEnum(str, Enum): + 'accepts only string values' + class Name(StrEnum): + BDFL = 'Guido van Rossum' + FLUFL = 'Barry Warsaw' +except Exception as exc: + Name = exc + +try: + Question = Enum('Question', 'who what when where why', module=__name__) +except Exception as exc: + Question = exc + +try: + Answer = Enum('Answer', 'him this then there because') +except Exception as exc: + Answer = exc + +# for doctests +try: + class Fruit(Enum): + tomato = 1 + banana = 2 + cherry = 3 +except Exception: + pass + +class TestEnum(unittest.TestCase): + def setUp(self): + class Season(Enum): + SPRING = 1 + SUMMER = 2 + AUTUMN = 3 + WINTER = 4 + self.Season = Season + + def test_enum_in_enum_out(self): + Season = self.Season + self.assertIs(Season(Season.WINTER), Season.WINTER) + + def test_enum_value(self): + Season = self.Season + self.assertEqual(Season.SPRING.value, 1) + + def test_intenum_value(self): + self.assertEqual(IntStooges.CURLY.value, 2) + + def test_dir_on_class(self): + Season = self.Season + self.assertEqual( + set(dir(Season)), + set(['__class__', '__doc__', '__members__', + 'SPRING', 'SUMMER', 'AUTUMN', 'WINTER']), + ) + + def test_dir_on_item(self): + Season = self.Season + self.assertEqual( + set(dir(Season.WINTER)), + set(['__class__', '__doc__', 'name', 'value']), + ) + + def test_enum(self): + Season = self.Season + lst = list(Season) + self.assertEqual(len(lst), len(Season)) + self.assertEqual(len(Season), 4, Season) + self.assertEqual( + [Season.SPRING, Season.SUMMER, Season.AUTUMN, Season.WINTER], lst) + + for i, season in enumerate('SPRING SUMMER AUTUMN WINTER'.split(), 1): + e = Season(i) + self.assertEqual(e, getattr(Season, season)) + self.assertEqual(e.value, i) + self.assertNotEqual(e, i) + self.assertEqual(e.name, season) + self.assertIn(e, Season) + self.assertIs(type(e), Season) + self.assertIsInstance(e, Season) + self.assertEqual(str(e), 'Season.' + season) + self.assertEqual( + repr(e), + '<Season.{0}: {1}>'.format(season, i), + ) + + def test_value_name(self): + Season = self.Season + self.assertEqual(Season.SPRING.name, 'SPRING') + self.assertEqual(Season.SPRING.value, 1) + with self.assertRaises(AttributeError): + Season.SPRING.name = 'invierno' + with self.assertRaises(AttributeError): + Season.SPRING.value = 2 + + def test_invalid_names(self): + with self.assertRaises(ValueError): + class Wrong(Enum): + mro = 9 + with self.assertRaises(ValueError): + class Wrong(Enum): + _create_= 11 + with self.assertRaises(ValueError): + class Wrong(Enum): + _get_mixins_ = 9 + with self.assertRaises(ValueError): + class Wrong(Enum): + _find_new_ = 1 + with self.assertRaises(ValueError): + class Wrong(Enum): + _any_name_ = 9 + + def test_contains(self): + Season = self.Season + self.assertIn(Season.AUTUMN, Season) + self.assertNotIn(3, Season) + + val = Season(3) + self.assertIn(val, Season) + + class OtherEnum(Enum): + one = 1; two = 2 + self.assertNotIn(OtherEnum.two, Season) + + def test_comparisons(self): + Season = self.Season + with self.assertRaises(TypeError): + Season.SPRING < Season.WINTER + with self.assertRaises(TypeError): + Season.SPRING > 4 + + self.assertNotEqual(Season.SPRING, 1) + + class Part(Enum): + SPRING = 1 + CLIP = 2 + BARREL = 3 + + self.assertNotEqual(Season.SPRING, Part.SPRING) + with self.assertRaises(TypeError): + Season.SPRING < Part.CLIP + + def test_enum_duplicates(self): + class Season(Enum): + SPRING = 1 + SUMMER = 2 + AUTUMN = FALL = 3 + WINTER = 4 + ANOTHER_SPRING = 1 + lst = list(Season) + self.assertEqual( + lst, + [Season.SPRING, Season.SUMMER, + Season.AUTUMN, Season.WINTER, + ]) + self.assertIs(Season.FALL, Season.AUTUMN) + self.assertEqual(Season.FALL.value, 3) + self.assertEqual(Season.AUTUMN.value, 3) + self.assertIs(Season(3), Season.AUTUMN) + self.assertIs(Season(1), Season.SPRING) + self.assertEqual(Season.FALL.name, 'AUTUMN') + self.assertEqual( + [k for k,v in Season.__members__.items() if v.name != k], + ['FALL', 'ANOTHER_SPRING'], + ) + + def test_enum_with_value_name(self): + class Huh(Enum): + name = 1 + value = 2 + self.assertEqual( + list(Huh), + [Huh.name, Huh.value], + ) + self.assertIs(type(Huh.name), Huh) + self.assertEqual(Huh.name.name, 'name') + self.assertEqual(Huh.name.value, 1) + def test_hash(self): + Season = self.Season + dates = {} + dates[Season.WINTER] = '1225' + dates[Season.SPRING] = '0315' + dates[Season.SUMMER] = '0704' + dates[Season.AUTUMN] = '1031' + self.assertEqual(dates[Season.AUTUMN], '1031') + + def test_intenum_from_scratch(self): + class phy(int, Enum): + pi = 3 + tau = 2 * pi + self.assertTrue(phy.pi < phy.tau) + + def test_intenum_inherited(self): + class IntEnum(int, Enum): + pass + class phy(IntEnum): + pi = 3 + tau = 2 * pi + self.assertTrue(phy.pi < phy.tau) + + def test_floatenum_from_scratch(self): + class phy(float, Enum): + pi = 3.141596 + tau = 2 * pi + self.assertTrue(phy.pi < phy.tau) + + def test_floatenum_inherited(self): + class FloatEnum(float, Enum): + pass + class phy(FloatEnum): + pi = 3.141596 + tau = 2 * pi + self.assertTrue(phy.pi < phy.tau) + + def test_strenum_from_scratch(self): + class phy(str, Enum): + pi = 'Pi' + tau = 'Tau' + self.assertTrue(phy.pi < phy.tau) + + def test_strenum_inherited(self): + class StrEnum(str, Enum): + pass + class phy(StrEnum): + pi = 'Pi' + tau = 'Tau' + self.assertTrue(phy.pi < phy.tau) + + + def test_intenum(self): + class WeekDay(IntEnum): + SUNDAY = 1 + MONDAY = 2 + TUESDAY = 3 + WEDNESDAY = 4 + THURSDAY = 5 + FRIDAY = 6 + SATURDAY = 7 + + self.assertEqual(['a', 'b', 'c'][WeekDay.MONDAY], 'c') + self.assertEqual([i for i in range(WeekDay.TUESDAY)], [0, 1, 2]) + + lst = list(WeekDay) + self.assertEqual(len(lst), len(WeekDay)) + self.assertEqual(len(WeekDay), 7) + target = 'SUNDAY MONDAY TUESDAY WEDNESDAY THURSDAY FRIDAY SATURDAY' + target = target.split() + for i, weekday in enumerate(target, 1): + e = WeekDay(i) + self.assertEqual(e, i) + self.assertEqual(int(e), i) + self.assertEqual(e.name, weekday) + self.assertIn(e, WeekDay) + self.assertEqual(lst.index(e)+1, i) + self.assertTrue(0 < e < 8) + self.assertIs(type(e), WeekDay) + self.assertIsInstance(e, int) + self.assertIsInstance(e, Enum) + + def test_intenum_duplicates(self): + class WeekDay(IntEnum): + SUNDAY = 1 + MONDAY = 2 + TUESDAY = TEUSDAY = 3 + WEDNESDAY = 4 + THURSDAY = 5 + FRIDAY = 6 + SATURDAY = 7 + self.assertIs(WeekDay.TEUSDAY, WeekDay.TUESDAY) + self.assertEqual(WeekDay(3).name, 'TUESDAY') + self.assertEqual([k for k,v in WeekDay.__members__.items() + if v.name != k], ['TEUSDAY', ]) + + def test_pickle_enum(self): + if isinstance(Stooges, Exception): + raise Stooges + self.assertIs(Stooges.CURLY, loads(dumps(Stooges.CURLY))) + self.assertIs(Stooges, loads(dumps(Stooges))) + + def test_pickle_int(self): + if isinstance(IntStooges, Exception): + raise IntStooges + self.assertIs(IntStooges.CURLY, loads(dumps(IntStooges.CURLY))) + self.assertIs(IntStooges, loads(dumps(IntStooges))) + + def test_pickle_float(self): + if isinstance(FloatStooges, Exception): + raise FloatStooges + self.assertIs(FloatStooges.CURLY, loads(dumps(FloatStooges.CURLY))) + self.assertIs(FloatStooges, loads(dumps(FloatStooges))) + + def test_pickle_enum_function(self): + if isinstance(Answer, Exception): + raise Answer + self.assertIs(Answer.him, loads(dumps(Answer.him))) + self.assertIs(Answer, loads(dumps(Answer))) + + def test_pickle_enum_function_with_module(self): + if isinstance(Question, Exception): + raise Question + self.assertIs(Question.who, loads(dumps(Question.who))) + self.assertIs(Question, loads(dumps(Question))) + + def test_exploding_pickle(self): + BadPickle = Enum('BadPickle', 'dill sweet bread-n-butter') + enum._make_class_unpicklable(BadPickle) + globals()['BadPickle'] = BadPickle + with self.assertRaises(TypeError): + dumps(BadPickle.dill) + with self.assertRaises(PicklingError): + dumps(BadPickle) + + def test_string_enum(self): + class SkillLevel(str, Enum): + master = 'what is the sound of one hand clapping?' + journeyman = 'why did the chicken cross the road?' + apprentice = 'knock, knock!' + self.assertEqual(SkillLevel.apprentice, 'knock, knock!') + + def test_getattr_getitem(self): + class Period(Enum): + morning = 1 + noon = 2 + evening = 3 + night = 4 + self.assertIs(Period(2), Period.noon) + self.assertIs(getattr(Period, 'night'), Period.night) + self.assertIs(Period['morning'], Period.morning) + + def test_getattr_dunder(self): + Season = self.Season + self.assertTrue(getattr(Season, '__eq__')) + + def test_iteration_order(self): + class Season(Enum): + SUMMER = 2 + WINTER = 4 + AUTUMN = 3 + SPRING = 1 + self.assertEqual( + list(Season), + [Season.SUMMER, Season.WINTER, Season.AUTUMN, Season.SPRING], + ) + + def test_programatic_function_string(self): + SummerMonth = Enum('SummerMonth', 'june july august') + lst = list(SummerMonth) + self.assertEqual(len(lst), len(SummerMonth)) + self.assertEqual(len(SummerMonth), 3, SummerMonth) + self.assertEqual( + [SummerMonth.june, SummerMonth.july, SummerMonth.august], + lst, + ) + for i, month in enumerate('june july august'.split(), 1): + e = SummerMonth(i) + self.assertEqual(int(e.value), i) + self.assertNotEqual(e, i) + self.assertEqual(e.name, month) + self.assertIn(e, SummerMonth) + self.assertIs(type(e), SummerMonth) + + def test_programatic_function_string_list(self): + SummerMonth = Enum('SummerMonth', ['june', 'july', 'august']) + lst = list(SummerMonth) + self.assertEqual(len(lst), len(SummerMonth)) + self.assertEqual(len(SummerMonth), 3, SummerMonth) + self.assertEqual( + [SummerMonth.june, SummerMonth.july, SummerMonth.august], + lst, + ) + for i, month in enumerate('june july august'.split(), 1): + e = SummerMonth(i) + self.assertEqual(int(e.value), i) + self.assertNotEqual(e, i) + self.assertEqual(e.name, month) + self.assertIn(e, SummerMonth) + self.assertIs(type(e), SummerMonth) + + def test_programatic_function_iterable(self): + SummerMonth = Enum( + 'SummerMonth', + (('june', 1), ('july', 2), ('august', 3)) + ) + lst = list(SummerMonth) + self.assertEqual(len(lst), len(SummerMonth)) + self.assertEqual(len(SummerMonth), 3, SummerMonth) + self.assertEqual( + [SummerMonth.june, SummerMonth.july, SummerMonth.august], + lst, + ) + for i, month in enumerate('june july august'.split(), 1): + e = SummerMonth(i) + self.assertEqual(int(e.value), i) + self.assertNotEqual(e, i) + self.assertEqual(e.name, month) + self.assertIn(e, SummerMonth) + self.assertIs(type(e), SummerMonth) + + def test_programatic_function_from_dict(self): + SummerMonth = Enum( + 'SummerMonth', + OrderedDict((('june', 1), ('july', 2), ('august', 3))) + ) + lst = list(SummerMonth) + self.assertEqual(len(lst), len(SummerMonth)) + self.assertEqual(len(SummerMonth), 3, SummerMonth) + self.assertEqual( + [SummerMonth.june, SummerMonth.july, SummerMonth.august], + lst, + ) + for i, month in enumerate('june july august'.split(), 1): + e = SummerMonth(i) + self.assertEqual(int(e.value), i) + self.assertNotEqual(e, i) + self.assertEqual(e.name, month) + self.assertIn(e, SummerMonth) + self.assertIs(type(e), SummerMonth) + + def test_programatic_function_type(self): + SummerMonth = Enum('SummerMonth', 'june july august', type=int) + lst = list(SummerMonth) + self.assertEqual(len(lst), len(SummerMonth)) + self.assertEqual(len(SummerMonth), 3, SummerMonth) + self.assertEqual( + [SummerMonth.june, SummerMonth.july, SummerMonth.august], + lst, + ) + for i, month in enumerate('june july august'.split(), 1): + e = SummerMonth(i) + self.assertEqual(e, i) + self.assertEqual(e.name, month) + self.assertIn(e, SummerMonth) + self.assertIs(type(e), SummerMonth) + + def test_programatic_function_type_from_subclass(self): + SummerMonth = IntEnum('SummerMonth', 'june july august') + lst = list(SummerMonth) + self.assertEqual(len(lst), len(SummerMonth)) + self.assertEqual(len(SummerMonth), 3, SummerMonth) + self.assertEqual( + [SummerMonth.june, SummerMonth.july, SummerMonth.august], + lst, + ) + for i, month in enumerate('june july august'.split(), 1): + e = SummerMonth(i) + self.assertEqual(e, i) + self.assertEqual(e.name, month) + self.assertIn(e, SummerMonth) + self.assertIs(type(e), SummerMonth) + + def test_subclassing(self): + if isinstance(Name, Exception): + raise Name + self.assertEqual(Name.BDFL, 'Guido van Rossum') + self.assertTrue(Name.BDFL, Name('Guido van Rossum')) + self.assertIs(Name.BDFL, getattr(Name, 'BDFL')) + self.assertIs(Name.BDFL, loads(dumps(Name.BDFL))) + + def test_extending(self): + class Color(Enum): + red = 1 + green = 2 + blue = 3 + with self.assertRaises(TypeError): + class MoreColor(Color): + cyan = 4 + magenta = 5 + yellow = 6 + + def test_exclude_methods(self): + class whatever(Enum): + this = 'that' + these = 'those' + def really(self): + return 'no, not %s' % self.value + self.assertIsNot(type(whatever.really), whatever) + self.assertEqual(whatever.this.really(), 'no, not that') + + def test_overwrite_enums(self): + class Why(Enum): + question = 1 + answer = 2 + propisition = 3 + def question(self): + print(42) + self.assertIsNot(type(Why.question), Why) + self.assertNotIn(Why.question, Why._member_names) + self.assertNotIn(Why.question, Why) + + def test_wrong_inheritance_order(self): + with self.assertRaises(TypeError): + class Wrong(Enum, str): + NotHere = 'error before this point' + + def test_intenum_transitivity(self): + class number(IntEnum): + one = 1 + two = 2 + three = 3 + class numero(IntEnum): + uno = 1 + dos = 2 + tres = 3 + self.assertEqual(number.one, numero.uno) + self.assertEqual(number.two, numero.dos) + self.assertEqual(number.three, numero.tres) + + def test_wrong_enum_in_call(self): + class Monochrome(Enum): + black = 0 + white = 1 + class Gender(Enum): + male = 0 + female = 1 + self.assertRaises(ValueError, Monochrome, Gender.male) + + def test_wrong_enum_in_mixed_call(self): + class Monochrome(IntEnum): + black = 0 + white = 1 + class Gender(Enum): + male = 0 + female = 1 + self.assertRaises(ValueError, Monochrome, Gender.male) + + def test_mixed_enum_in_call_1(self): + class Monochrome(IntEnum): + black = 0 + white = 1 + class Gender(IntEnum): + male = 0 + female = 1 + self.assertIs(Monochrome(Gender.female), Monochrome.white) + + def test_mixed_enum_in_call_2(self): + class Monochrome(Enum): + black = 0 + white = 1 + class Gender(IntEnum): + male = 0 + female = 1 + self.assertIs(Monochrome(Gender.male), Monochrome.black) + + def test_flufl_enum(self): + class Fluflnum(Enum): + def __int__(self): + return int(self.value) + class MailManOptions(Fluflnum): + option1 = 1 + option2 = 2 + option3 = 3 + self.assertEqual(int(MailManOptions.option1), 1) + + def test_no_such_enum_member(self): + class Color(Enum): + red = 1 + green = 2 + blue = 3 + with self.assertRaises(ValueError): + Color(4) + with self.assertRaises(KeyError): + Color['chartreuse'] + + def test_new_repr(self): + class Color(Enum): + red = 1 + green = 2 + blue = 3 + def __repr__(self): + return "don't you just love shades of %s?" % self.name + self.assertEqual( + repr(Color.blue), + "don't you just love shades of blue?", + ) + + def test_inherited_repr(self): + class MyEnum(Enum): + def __repr__(self): + return "My name is %s." % self.name + class MyIntEnum(int, MyEnum): + this = 1 + that = 2 + theother = 3 + self.assertEqual(repr(MyIntEnum.that), "My name is that.") + + def test_multiple_mixin_mro(self): + class auto_enum(type(Enum)): + def __new__(metacls, cls, bases, classdict): + temp = type(classdict)() + names = set(classdict._member_names) + i = 0 + for k in classdict._member_names: + v = classdict[k] + if v is Ellipsis: + v = i + else: + i = v + i += 1 + temp[k] = v + for k, v in classdict.items(): + if k not in names: + temp[k] = v + return super(auto_enum, metacls).__new__( + metacls, cls, bases, temp) + + class AutoNumberedEnum(Enum, metaclass=auto_enum): + pass + + class AutoIntEnum(IntEnum, metaclass=auto_enum): + pass + + class TestAutoNumber(AutoNumberedEnum): + a = ... + b = 3 + c = ... + + class TestAutoInt(AutoIntEnum): + a = ... + b = 3 + c = ... + + def test_subclasses_with_getnewargs(self): + class NamedInt(int): + def __new__(cls, *args): + _args = args + name, *args = args + if len(args) == 0: + raise TypeError("name and value must be specified") + self = int.__new__(cls, *args) + self._intname = name + self._args = _args + return self + def __getnewargs__(self): + return self._args + @property + def __name__(self): + return self._intname + def __repr__(self): + # repr() is updated to include the name and type info + return "{}({!r}, {})".format(type(self).__name__, + self.__name__, + int.__repr__(self)) + def __str__(self): + # str() is unchanged, even if it relies on the repr() fallback + base = int + base_str = base.__str__ + if base_str.__objclass__ is object: + return base.__repr__(self) + return base_str(self) + # for simplicity, we only define one operator that + # propagates expressions + def __add__(self, other): + temp = int(self) + int( other) + if isinstance(self, NamedInt) and isinstance(other, NamedInt): + return NamedInt( + '({0} + {1})'.format(self.__name__, other.__name__), + temp ) + else: + return temp + + class NEI(NamedInt, Enum): + x = ('the-x', 1) + y = ('the-y', 2) + + self.assertIs(NEI.__new__, Enum.__new__) + self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)") + globals()['NamedInt'] = NamedInt + globals()['NEI'] = NEI + NI5 = NamedInt('test', 5) + self.assertEqual(NI5, 5) + self.assertEqual(loads(dumps(NI5)), 5) + self.assertEqual(NEI.y.value, 2) + self.assertIs(loads(dumps(NEI.y)), NEI.y) + + def test_subclasses_without_getnewargs(self): + class NamedInt(int): + def __new__(cls, *args): + _args = args + name, *args = args + if len(args) == 0: + raise TypeError("name and value must be specified") + self = int.__new__(cls, *args) + self._intname = name + self._args = _args + return self + @property + def __name__(self): + return self._intname + def __repr__(self): + # repr() is updated to include the name and type info + return "{}({!r}, {})".format(type(self).__name__, + self.__name__, + int.__repr__(self)) + def __str__(self): + # str() is unchanged, even if it relies on the repr() fallback + base = int + base_str = base.__str__ + if base_str.__objclass__ is object: + return base.__repr__(self) + return base_str(self) + # for simplicity, we only define one operator that + # propagates expressions + def __add__(self, other): + temp = int(self) + int( other) + if isinstance(self, NamedInt) and isinstance(other, NamedInt): + return NamedInt( + '({0} + {1})'.format(self.__name__, other.__name__), + temp ) + else: + return temp + + class NEI(NamedInt, Enum): + x = ('the-x', 1) + y = ('the-y', 2) + + self.assertIs(NEI.__new__, Enum.__new__) + self.assertEqual(repr(NEI.x + NEI.y), "NamedInt('(the-x + the-y)', 3)") + globals()['NamedInt'] = NamedInt + globals()['NEI'] = NEI + NI5 = NamedInt('test', 5) + self.assertEqual(NI5, 5) + self.assertEqual(NEI.y.value, 2) + with self.assertRaises(TypeError): + dumps(NEI.x) + with self.assertRaises(PicklingError): + dumps(NEI) + + def test_tuple_subclass(self): + class SomeTuple(tuple, Enum): + first = (1, 'for the money') + second = (2, 'for the show') + third = (3, 'for the music') + self.assertIs(type(SomeTuple.first), SomeTuple) + self.assertIsInstance(SomeTuple.second, tuple) + self.assertEqual(SomeTuple.third, (3, 'for the music')) + globals()['SomeTuple'] = SomeTuple + self.assertIs(loads(dumps(SomeTuple.first)), SomeTuple.first) + + def test_duplicate_values_give_unique_enum_items(self): + class AutoNumber(Enum): + first = () + second = () + third = () + def __new__(cls): + value = len(cls.__members__) + 1 + obj = object.__new__(cls) + obj._value = value + return obj + def __int__(self): + return int(self._value) + self.assertEqual( + list(AutoNumber), + [AutoNumber.first, AutoNumber.second, AutoNumber.third], + ) + self.assertEqual(int(AutoNumber.second), 2) + self.assertIs(AutoNumber(1), AutoNumber.first) + + def test_inherited_new_from_enhanced_enum(self): + class AutoNumber(Enum): + def __new__(cls): + value = len(cls.__members__) + 1 + obj = object.__new__(cls) + obj._value = value + return obj + def __int__(self): + return int(self._value) + class Color(AutoNumber): + red = () + green = () + blue = () + self.assertEqual(list(Color), [Color.red, Color.green, Color.blue]) + self.assertEqual(list(map(int, Color)), [1, 2, 3]) + + def test_inherited_new_from_mixed_enum(self): + class AutoNumber(IntEnum): + def __new__(cls): + value = len(cls.__members__) + 1 + obj = int.__new__(cls, value) + obj._value = value + return obj + class Color(AutoNumber): + red = () + green = () + blue = () + self.assertEqual(list(Color), [Color.red, Color.green, Color.blue]) + self.assertEqual(list(map(int, Color)), [1, 2, 3]) + + def test_ordered_mixin(self): + class OrderedEnum(Enum): + def __ge__(self, other): + if self.__class__ is other.__class__: + return self._value >= other._value + return NotImplemented + def __gt__(self, other): + if self.__class__ is other.__class__: + return self._value > other._value + return NotImplemented + def __le__(self, other): + if self.__class__ is other.__class__: + return self._value <= other._value + return NotImplemented + def __lt__(self, other): + if self.__class__ is other.__class__: + return self._value < other._value + return NotImplemented + class Grade(OrderedEnum): + A = 5 + B = 4 + C = 3 + D = 2 + F = 1 + self.assertGreater(Grade.A, Grade.B) + self.assertLessEqual(Grade.F, Grade.C) + self.assertLess(Grade.D, Grade.A) + self.assertGreaterEqual(Grade.B, Grade.B) + def test_extending2(self): + class Shade(Enum): + def shade(self): + print(self.name) + class Color(Shade): + red = 1 + green = 2 + blue = 3 + with self.assertRaises(TypeError): + class MoreColor(Color): + cyan = 4 + magenta = 5 + yellow = 6 + + def test_extending3(self): + class Shade(Enum): + def shade(self): + return self.name + class Color(Shade): + def hex(self): + return '%s hexlified!' % self.value + class MoreColor(Color): + cyan = 4 + magenta = 5 + yellow = 6 + self.assertEqual(MoreColor.magenta.hex(), '5 hexlified!') + + + def test_no_duplicates(self): + class UniqueEnum(Enum): + def __init__(self, *args): + cls = self.__class__ + if any(self.value == e.value for e in cls): + a = self.name + e = cls(self.value).name + raise ValueError( + "aliases not allowed in UniqueEnum: %r --> %r" + % (a, e) + ) + class Color(UniqueEnum): + red = 1 + green = 2 + blue = 3 + with self.assertRaises(ValueError): + class Color(UniqueEnum): + red = 1 + green = 2 + blue = 3 + grene = 2 + + def test_init(self): + class Planet(Enum): + MERCURY = (3.303e+23, 2.4397e6) + VENUS = (4.869e+24, 6.0518e6) + EARTH = (5.976e+24, 6.37814e6) + MARS = (6.421e+23, 3.3972e6) + JUPITER = (1.9e+27, 7.1492e7) + SATURN = (5.688e+26, 6.0268e7) + URANUS = (8.686e+25, 2.5559e7) + NEPTUNE = (1.024e+26, 2.4746e7) + def __init__(self, mass, radius): + self.mass = mass # in kilograms + self.radius = radius # in meters + @property + def surface_gravity(self): + # universal gravitational constant (m3 kg-1 s-2) + G = 6.67300E-11 + return G * self.mass / (self.radius * self.radius) + self.assertEqual(round(Planet.EARTH.surface_gravity, 2), 9.80) + self.assertEqual(Planet.EARTH.value, (5.976e+24, 6.37814e6)) + + +if __name__ == '__main__': + unittest.main() @@ -381,6 +381,8 @@ Library - Implement PEP 443 "Single-dispatch generic functions". +- Implement PEP 435 "Adding an Enum type to the Python standard library". + Tests ----- |