The quest for the perfect python enum/constant declaration, Part 2

# constant.py -
# helper classes for enumerating constants of ints and bit-based flags.
#
# Copyright (C) 2008, Aldrin Martoq
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

__date__ = "2008-12-22"
__author__ = "Aldrin Martoq "
__license__ = "LGPL"
__version__ = "1.0.0"
__copyright__ = "Copyright (C) 2008 Aldrin Martoq"
__URL__ = "http://aldrin.martoq.cl/techblog/"
__all__ = "ConstantInt ConstantBit ConstantEnum ConstantFlags".split()

class ConstantInt(int):
""" An int with a name associated to a ConstEnum. Yea! """

def __new__(self, value, name, enum):
return int.__new__(self, value)
def __init__(self, value, name, enum):
self._name = name
self._enum = enum
def _convert(self, other):
return self._enum[other]
## numeric support ##
def __add__(self, other):
return self._convert(int.__add__(self, other))
def __sub__(self, other):
return self._convert(int.__sub__(self, other))
def __mul__(self, other):
return self._convert(int.__mul__(self, other))
def __divmod__(self, other):
return self._convert(int.__divmod__(self, other))
def __pow__(self, other):
return self._convert(int.__pow__(self, other))
def __lshift__(self, other):
return self._convert(int.__lshift__(self, other))
def __rshift__(self, other):
return self._convert(int.__rshift__(self, other))
def __and__(self, other):
return self._convert(int.__and__(self, other))
def __xor__(self, other):
return self._convert(int.__xor__(self, other))
def __or__(self, other):
return self._convert(int.__or__(self, other))
def __div__(self, other):
return self._convert(int.__div__(self, other))
def __truediv__(self, other):
return self._convert(int.__truediv__(self, other))
def __neg__(self, other):
return self._convert(int.__neg__(self, other))
def __pos__(self, other):
return self._convert(int.__pos__(self, other))
def __abs__(self, other):
return self._convert(int.__abs__(self, other))
def __invert__(self, other):
return self._convert(int.__invert__(self, other))

## eq support ##
def __eq__(self, other):
if isinstance(other, str):
return self._name == other
return int(self) == other
def __ne__(self, other):
return not self == other

## string representation ##
def __str__(self):
return "%s" % (self._name)
def __repr__(self):
return "%s(%d)" % (self._name, self)

class ConstantBit(ConstantInt):
""" An int treated like a bit flag. Yea! """
def _mask(self):
r = []
d = int(self)
for v in self._enum:
if int(v) & d or v == d:
r.append(ConstantInt.__str__(v))
d -= int(v)
if len(r) == 0 or d != 0: r.append(str(d))
return r
def __str__(self):
return "|".join(self._mask())
def __repr__(self):
return "%s(0x%x)" % ("|".join(self._mask()), self)

class ConstantEnum(object):
""" Container for named integers. Yea!

Sample Usage:
>>> NUMBERS = ConstantEnum("UNO DOS SIETE=7 OCHO")
>>> NUMBERS

>>> str(NUMBERS.UNO), repr(NUMBERS.UNO), int(NUMBERS.UNO)
('UNO', 'UNO(1)', 1)
>>> NUMBERS.UNO + 1, NUMBERS.UNO - 1
(DOS(2), 0(0))
>>> NUMBERS[7], NUMBERS["UNO"]
(SIETE(7), UNO(1))
>>> [i for i in NUMBERS]
[OCHO(8), UNO(1), DOS(2), SIETE(7)]

"""
def __init__(self, items, atype = ConstantInt, start_value = 1):
self._atype = atype # type for creating instances of this enum
self._names = {} # instances by name
self._values = {} # instances by value
# now we put the items
value = start_value
for item in items.split():
s = item.split("=")
if len(s) == 2: value = int(s[1])
self.put(value, s[0])
value = self._inc(value)
## internal support ##
def _inc(self, value):
return value + 1
def _create(self, value, name = None):
return self._atype(value, name, self)
def _find_or_create(self, value):
try:
if isinstance(value, str):
if self._names.has_key(value):
return self._names[value]
else:
return None
return self._values[value]
except:
return self._create(value, str(value))
## external api ##
def put(self, value, name):
""" creates a new constant given a value and name """
newobj = self._create(value, name)
self._names[name] = newobj
self._values[value] = newobj
setattr(self, name, newobj)
def remove(self, name_or_value):
""" removes an existing constant either by name or value """
if isinstance(name_or_value, str):
obj = self._names[name_or_value]
else:
obj = self._values[name_or_value]
del self._names[str(obj)]
del self._values[int(obj)]
## iter and item support ##
def __iter__(self):
return iter(self._values.values())
def __len__(self):
return len(self._values)
def __getitem__(self, value):
return self._find_or_create(value)
## str related ##
def __str__(self):
return "<%s %s %s>" % (self.__class__.__name__,
self._atype.__name__,
[i._name for i in self])
def __repr__(self):
return "<%s %s %s>" % (self.__class__.__name__,
self._atype.__name__,
["%s(%d)" % (i._name, i) for i in self])

class ConstantFlags(ConstantEnum):
""" Container for bit-based flags. Yea!

Sample Usage:
>>> FLAGS = ConstantFlags("RED GREEN BLUE NONE=0")
>>> FLAGS

>>> FLAGS.RED, FLAGS.GREEN, FLAGS.RED|FLAGS.GREEN, FLAGS[3]
(RED(0x1), GREEN(0x2), RED|GREEN(0x3), RED|GREEN(0x3))
>>> FLAGS[0], FLAGS[1], FLAGS["NONE"], FLAGS["BLUE"]
(NONE(0x0), RED(0x1), NONE(0x0), BLUE(0x4))
>>> [i for i in FLAGS]
[NONE(0x0), RED(0x1), GREEN(0x2), BLUE(0x4)]

"""
def __init__(self, items, atype = ConstantBit, start_value = 1):
ConstantEnum.__init__(self, items, atype, start_value)
def _inc(self, value):
return value * 2