Welcome, guest | Sign In | My Account | Store | Cart

Yet another enum for Python (Python recipe) by Antonio Cuni
ActiveState Code (http://code.activestate.com/recipes/576711/)

There are plenty of ways to declare C-like enums in Python. This one tries to be minimalistic and with a nice syntax, (ab)using decorators.

The code is also available here: http://codespeak.net/svn/user/antocuni/hack/enum.py

Python, 55 lines
Download
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class EnumItem(int):
    
    __slots__ = ['name']

    def __new__(cls, name, value):
        i = int.__new__(cls, value)
        i.name = name
        return i

    def getname(self):
        return self.name


def enum(func):
    names = func.func_code.co_varnames
    defaults = func.func_defaults
    if defaults is None:
        defaults = []
    n = len(names)-len(defaults)
    values = range(n)+list(defaults)
    for i, name in zip(values, names):
        item = EnumItem(name, i)
        setattr(func, name, item)
    return func




# ________________________________________________
# tests

def test_enum():
    @enum
    def colors(red, green, blue):
        pass
    assert colors.red == 0
    assert colors.green == 1
    assert colors.blue == 2


def test_getname():
    @enum
    def colors(red, green, blue):
        pass
    assert colors.red.getname() == 'red'
    assert colors.green.getname() == 'green'
    assert colors.blue.getname() == 'blue'

def test_default_value():
    @enum
    def colors(red, green, blue=42):
        pass
    assert colors.red == 0
    assert colors.green == 1
    assert colors.blue == 42
Tags: enum

2 comments

Richard Nichols III 15 years ago  # | flag

colors = enum(lambda red, green, blue: None)

:)

with a little modification, namely changing "names = func.func_code.co_varnames" to "names = func.func_code.co_names", I think we can do this..:

colors = enum(lambda: (red, green, blue))

(not that I know that anybody would want to. =P)

Matteo Dell'Amico 15 years ago  # | flag

Hey Anto ;)

I think that having to use strings wouldn't be a big deal, if the names gets separated in the way namedtuple does:

from collections import namedtuple

def enum(*args):
    t = namedtuple("Enum", *args)
    return t._make(range(len(t._fields)))

This way you get

>>> colors = enum('red green blue') # or enum('red, green, blue')
>>> colors.red
0
>>> colors.green
1
>>> colors.blue
2
>>> colors
Enum(red=0, green=1, blue=2)

If you really want to change the default values you could play around with a **kw argument, but I don't see use cases for this...

Created by Antonio Cuni on Fri, 3 Apr 2009 (MIT)
◄ Python recipes (4591) ►
◄ Antonio Cuni's recipes (2) ►

Required Modules

  • (none specified)

Other Information and Tasks