enumeration.py

"""
Helpers to define enumeration classes.

An enumeration is a "namespace" or "class" of constants that define a number
of values within that namespace that are canonical or descriptive, and usually
represent the full range of possible values within that realm. This allows for
having more descriptive and clear code and also helps prevent errors due to
mis-spelling of constants.

Simple Example::

    class Action(Enumeration):
        CREATE = 'create'
        EDIT   = 'edit'
        DELETE = 'delete'

get the value of individual constants:

    >>> Action.CREATE
    'create'
    >>> Action.EDIT
    'edit'

Convenience accessors for canonical lists:

    >>> Action.names()
    ['CREATE', 'EDIT', 'DELETE']
    >>> Action.values()
    ['create', 'edit', 'delete']
"""
__all__ = ('Enumeration', 'EnumAtom')


class EnumAtom(object):
    creation_counter = 0

    def __init__(self, *args, **kw):
        self.creation_counter = EnumAtom.creation_counter = EnumAtom.creation_counter + 1
        self.args = args
        self.kw = kw


def _initialize_values(d):
    """Helper to extract values"""
    if '_values' in d:
        return tuple(d.pop('_values'))
    else:
        _values = []
        has_atoms = True
        for name in d:
            if name.startswith('_'):
                continue
            v = d[name]
            if not isinstance(v, EnumAtom):
                has_atoms = False
            _values.append((name, v))

        if has_atoms:
            _fieldnames = d.get('_fieldnames', ('value', 'label', 'key'))
            from collections import namedtuple
            T = namedtuple('Atom', _fieldnames)
            _values.sort(key=lambda x: x[1].creation_counter)
            new_values = []
            for key, atom in _values:
                kw = atom.kw
                if 'key' in _fieldnames:
                    kw['key'] = key
                obj = T(*atom.args, **kw)
                new_values.append((key, obj))
                if 'value' in _fieldnames:
                    d[key] = obj.value
                else:
                    d[key] = obj
            return tuple(new_values)

        return tuple(_values)


class EnumerationMeta(type):
    """
    The meta-class of all enumerations.
    """
    def __new__(cls, name, bases, d):
        values = _initialize_values(d)
        d['_values'] = values
        d['_lookup'] = dict(values)
        return type.__new__(cls, name, bases, d)

    def __iter__(cls):
        return iter(cls._values)

    def __dir__(cls):
        return cls._lookup.keys()

    def __call__(cls, *args, **kw):
        raise TypeError("You cannot instantiate an enumeration.")

    def __setattr__(cls, name, value):
        raise AttributeError("Enumeration '%s' cannot have values added to it at run-time." % cls)


class Enumeration(object):
    """
    An enumeration base class.

    This can be used to make new enumeration classes via subclassing.
    """
    __metaclass__ = EnumerationMeta

    @classmethod
    def lookup(cls, field, value):
        """
        Retrieve an EnumAtom by looking up one of the fields.
        """
        for _, v in cls._values:
            if getattr(v, field) == value:
                return v

    @classmethod
    def names(cls):
        """
        Get the name of all the values in here, in definition order if applicable.
        """
        return list(x[0] for x in cls._values)

    @classmethod
    def values(cls):
        """
        Get all values in here, in definition order if applicable.
        """
        return list(x[1] for x in cls._values)

    @classmethod
    def pairs(cls):
        return cls._values

    @classmethod
    def wtforms_choices(cls):
        return list((x[1].value, x[1].label) for x in cls._values)

usage.py

### Simple enumeration

class Metric(Enumeration):
  STEPS = 'steps'
  WORKOUTS = 'workouts'
  
>>> Metric.STEPS
'steps'
>>> Metric.values()
['steps', 'workouts']


### Enumeration that can also keep other data

class Metric(Enumeration):
  STEPS = Atom('steps', 'Step Count')
  WORKOUTS = Atom('workouts', 'Num Workouts')

>>> Metric.wtforms_choices()
[('steps', 'Step Count'), ('workouts', 'Num Workouts')]

### Custom attributes

class Metric(Enumeration):
  _fieldnames = ('value', 'label', 'unit', 'key')
  
  STEPS = Atom('steps', 'Step Count', 'step')
  WORKOUTS = Atom('workouts', 'Num Workouts', unit='workout')
  WEIGHT = Atom('weight', 'Weight Loss', unit='pounds')

>>> Metric.values()[-1]
Atom(value='weight', label='Weight Loss', unit='pounds', key='WEIGHT')
>>> Metric.values()[-1].unit
'pounds'