Wednesday, May 4, 2011

Enums in Python: a more flexible and powerful implementation

There is no standard implementation of enums (enumerated values) in Python. But enums are very useful and exist in most other programming languages. Hence several "recipes" for enum implementations have been proposed. As is clear from the differences between these implementations (and the comments that others have made about them), there are many different ideas about what the characteristics of an enum should be.

I needed enums for a Python project where the enums would be part of an API that I was defining for use by end-users. My requirements for an enum implementation were:

  1. Allow definition of a named class representing the set of allowable values. (I needed a class, not an instance, since I wanted to be able to specify the enum class as the parameter type for API functions that expected a class.) The enum classes should be derived from a common superclass (so I could test if a type was an enum type).
  2. The allowable values would be represented by instances of the enum class that would be automatically created at definition time. Each instance would have a name and a value (i.e. would be a name/value pair).
  3. The values should not be restricted to be integers (e.g. the values could be floats). But the values should be immutable objects (e.g. not lists)).
  4. Each enum class can be queried to get the list of allowable names and values.
  5. Each enum class can be queried to find the name/value pair instance given the name.
  6. It shouldn't be possible to modify the value of an existing name/value pair.
  7. Allow definition of alternative names for existing enum values. (e.g. so that the name can be changed in a future revision of the API while maintaining backward compatibility for users' source code) but these alternative names should not show up when the list of allowable values is queried.
None of the existing enum implementations fulfilled the above requirements, so I wrote my own. Here's the source code for my implementation: enum.py
(It is published under the MIT licence.)

To define an enum class, you use the 'Enum' function from the above 'enum' module. This function expects its first argument to be a string specifying the name of the enum class you are defining. The subsequent arguments should be keyword arguments specifying the names and values for this enum class. The 'Enum' function returns a class object (which is why its name is capitalized). For example, to define an enum class named 'Colours', you could do this:
from enum import Enum

>>> Colours = Enum('Colours', red=1, green=2, blue=3)

After that, there would be instances Colours.red, Colours.green, Colours.blue representing the possible colour values. To get the value of the 'red' instance, you would use 'Colours.red.value' (which is 1). To get the 'red' instance given a colour string in the variable 'col', you would use 'Colours.get(col)' or 'Colours(col)'. (If that string variable contained anything other than "red", "green" or "blue", you would get an exception.)

Some more examples of use will show the other facilities of this enum implementation:

from enum import Enum

>>> Colours = Enum('Colours', red=1, green=2, blue=3)
>>> Colours.green.name
'green'

>>> Colours.green.value
2

>>> Colours.names()
['red', 'green', 'blue']

>>> Colours.get('blue').value
3

>>> Colours.getByValue(3).name
'blue'

>>> Colours('blue').value
3

>>> Colours.addAlias('bleu', 'blue')
>>> Colours.names()
['red', 'green', 'blue']

>>> Colours.bleu.value
3

>>> [x.name for x in Colours.objects()]
['red', 'green', 'blue']

>>> listOfNames = ['red', 'blue']
>>> [x.name for x in Colours.listOfNamesToListOfEnums(listOfNames)]
['red', 'blue']

>>> commaSepStr = "red, blue"
>>> [x.name for x in Colours.commaSepStrToEnumList(commaSepStr)]
['red', 'blue']


>>> Numbers = Enum('Numbers', pi=3.1415926, e=2.71828)
>>> isinstance(Numbers, type)
True

>>> Numbers.names()
['e', 'pi']

>>> round(Numbers.e.value, 3)
2.718

>>> x = Numbers.pi
>>> x is Numbers('pi')
True

>>> Numbers('phi')
Traceback (most recent call last):
    ...
ValueError: 'phi' is not a valid enum name (Must be one of: ['e', 'pi'])

>>> Numbers.e.value = 42
Traceback (most recent call last):
    ...
TypeError: can't modify an enum

1 comment: