enum in Python

[Pettersen, Bjorn S]

> From: Scott David Daniels [mailto:Scott.Daniels at Acm.Org] 
> 
> David M. Cook wrote:
> 
> > In article <KF67b.33$fd2.25 at newssvr23.news.prodigy.com>, 
> > Andrew Chalk wrote:
> > 
> > 
> >>As a rank Python beginner I've used a dictionary, but 
> >>presumably there is a better way.
> > 
> > 
> > I've seen idioms like 
> > 
> > FOO, BAR, BAZ = range(3)
> > 
> > used.
> > 
> > Dave Cook
> 
> For 2.3 or after:
> 
>      class Enumerate(object):
>          def __init__(self, names):
>              for number, name in enumerate(names.split()):
>                  setattr(self, name, number)
> 
> To use:
>      codes = Enumerate('FOO BAR BAZ')
>      codes.BAZ will be 2 and so on.
[...]

Which is what I would use, however, sometimes you just want to reuse
your C code, so I had a little fun creating the following... Use as:

  from enum import enum
  class LoanTypeCode(enum): 
    """
        enum LOAN_TYPE_CODE {
            UNKNOWN = 0,
            CONSTRUCTION,
            FHA,
            VHA,
            LIVESTOCK = 100,
            MACHINERY_EQUIPMENT,
            FARM_REAL_ESTATE,
            CROPS_PROCEEDS_NEW,
            CROPS_PROCEEDS_OLD,
            DAIRY_PROCEEDS,
        };
    """

  >>> print LoanTypeCode
  enum LOAN_TYPE_CODE(UNKNOWN[0], CONSTRUCTION[1], FHA[2], VHA[3],
LIVESTOCK[100], 
  MACHINERY_EQUIPMENT[101], FARM_REAL_ESTATE[102],
CROPS_PROCEEDS_NEW[103], 
  CROPS_PROCEEDS_OLD[104], DAIRY_PROCEEDS[105])

Note the beautifully discrete values :-)

  >>> print LoanTypeCode.UNKNOWN
  <LOAN_TYPE_CODE: UNKNOWN[0]>
  
.. and enum values know both their name and numberic value (and almost
where they belong...) Almost always better to display UNKNOWN than 0..

-- bjorn

--- enum.py -----
def _docStringData(strRep):
    """Pull out relevant information from docstring... Hex values etc.
left as
       an exercise.
    """
    import re
    _name = ''
    _vals = []
    m = re.match(r'\s*enum\s*([\w_]+[\w\d_]*)\s*{\s*', strRep)
    if m:
        _name = m.group(1)
        strRep = strRep[m.end():]
        matches = re.findall(r'([\w_]+[\w\d_]*)(\s*=\s*(\d+))?', strRep)
        curval = 0
        for name, _, val in matches:
            if not val:
                curval += 1
            else:
                curval = int(val)
            _vals.append((name, curval))
        return (_name, _vals)
    
    raise ValueError("The input header doesn't look like an enum: %s" %
strRep)

class _enumItem(object):
    """Object to represent individual enum values."""
    __slots__ = ['_enumName', '_name', '_val']
    
    name = property(lambda self: self._name)
    val = property(lambda self: self._val)
    
    def __init__(self, enumName, name, value):
        self._enumName, self._name, self._val = enumName, name, value
    def __repr__(self):
        return '<%s: %s[%s]>' % (self._enumName, self.name,
str(self.val))
    def __str__(self):
        return '%s.%s' % (self._enumName, self.name)
    def __cmp__(self, other):
        return cmp(self.val, other.val)

class _enum(type):
    """Enum metaclass..."""
    def __new__(cls, name, bases, clsdict):
        if name in ('enum', '_enum'): # skip metaclass machinery
            return type.__new__(cls, name, bases, clsdict)
        
        if '__doc__' in clsdict:
            if '__init__' in clsdict:
                raise TypeError("Enums can't have __init__ methods")
                
            docstr = clsdict['__doc__']
            className, enumVals = _docStringData(docstr)
            _enumvals_ = {}
            clsdict['__enumvals__'] = _enumvals_
            for name, val in enumVals:
                item = _enumItem(className, name, val)
                _enumvals_[name] = item
                clsdict[name] = property(lambda
s,name=name:s.__enumvals__[name])
                
            klassdef = type.__new__(cls, className, bases, clsdict)
            return klassdef() # oops <wink>
            
        raise TypeError("class %s must have a doc string to be an enum"
% name)
    
class enum(object):
    __metaclass__ = _enum
    
    def __str__(self):
        return self.__class__.__name__
        
    def __repr__(self):
        header = "enum " + str(self) + '('
        # sort by values
        vals = self.__enumvals__.items()
        vals.sort(lambda x,y: cmp(x[1],y[1]))
        strvals = [name + '[' + str(val._val) + ']' for name, val in
vals]
        return header + ', '.join(strvals) + ')'