gistfile1.py

# Go-style enums in Python

#
# Implementation is here, scroll down for usage.
#

import operator

class LazyValue(object):
    def __init__(self, value):
        self.v = value

    __add__ = lambda s, o: LazyValue(lambda: s.v() + o)
    __sub__ = lambda s, o: LazyValue(lambda: s.v() - o)
    __mul__ = lambda s, o: LazyValue(lambda: s.v() * o)
    __truediv__ = lambda s, o: LazyValue(lambda: operator.truediv(s.v(), o))
    __floordiv__ = lambda s, o: LazyValue(lambda: s.v() // o)
    __mod__ = lambda s, o: LazyValue(lambda: s.v() % o)
    __divmod__ = lambda s, o: LazyValue(lambda: divmod(s.v(), o))
    __pow__ = lambda s, o, *a: LazyValue(lambda: pow(s.v(), o, *a))
    __lshift__ = lambda s, o: LazyValue(lambda: s.v() << o)
    __rshift__ = lambda s, o: LazyValue(lambda: s.v() >> o)
    __and__ = lambda s, o: LazyValue(lambda: s.v() & o)
    __xor__ = lambda s, o: LazyValue(lambda: s.v() ^ o)
    __or__ = lambda s, o: LazyValue(lambda: s.v() | o)

    __radd__ = lambda s, o: LazyValue(lambda: o + s.v())
    __rsub__ = lambda s, o: LazyValue(lambda: o - s.v())
    __rmul__ = lambda s, o: LazyValue(lambda: o * s.v())
    __rtruediv__ = lambda s, o: LazyValue(lambda: operator.truediv(o, s.v()))
    __rfloordiv__ = lambda s, o: LazyValue(lambda: o // s.v())
    __rmod__ = lambda s, o: LazyValue(lambda: o % s.v())
    __rdivmod__ = lambda s, o: LazyValue(lambda: divmod(o, s.v()))
    __rpow__ = lambda s, o, *a: LazyValue(lambda: pow(o, s.v(), *a))
    __rlshift__ = lambda s, o: LazyValue(lambda: o << s.v())
    __rrshift__ = lambda s, o: LazyValue(lambda: o >> s.v())
    __rand__ = lambda s, o: LazyValue(lambda: o & s.v())
    __rxor__ = lambda s, o: LazyValue(lambda: o ^ s.v())
    __ror__ = lambda s, o: LazyValue(lambda: o | s.v())

    __neg__ = lambda s: LazyValue(lambda: -s.v())
    __pos__ = lambda s: LazyValue(lambda: +s.v())
    __abs__ = lambda s: LazyValue(lambda: abs(s.v()))
    __invert__ = lambda s: LazyValue(lambda: ~s.v())

    __complex__ = lambda s: LazyValue(lambda: complex(s.v()))
    __int__ = lambda s: LazyValue(lambda: int(s.v()))
    __float__ = lambda s: LazyValue(lambda: float(s.v()))
    __round__ = lambda s, *a: LazyValue(lambda: round(s.v(), *a))

    __index__ = lambda s: LazyValue(lambda: s.v().__index__())

    def evaluate(self):
        value = self
        while isinstance(value, LazyValue):
            value = value.v()
        return value

class AutoIncrementDict(dict):
    def __init__(self, *args, **kwargs):
        self.last_iota = 0
        iota = LazyValue(lambda n=self.last_iota: n)
        self.in_expr = False
        self.current_iota_expr = iota
        self.current_iotas = [iota]
        super().__init__(*args, **kwargs)

    def __missing__(self, key):
        if key.startswith('__'):
            return
        for iota in self.current_iotas:
            iota.v = lambda n=self.last_iota: n
            self.last_iota += 1
        value = self.current_iota_expr.evaluate()
        super().__setitem__(key, value)

    def __getitem__(self, key):
        if key == 'iota':
            value = LazyValue(lambda n=self.last_iota: n)
            self.last_iota += 1
            if not self.in_expr:
                self.current_iotas = []
                self.in_expr = True
            self.current_iotas.append(value)
            return value
        return super().__getitem__(key)

    def __setitem__(self, key, value):
        if isinstance(value, int):
            self.last_iota = value + 1
        elif isinstance(value, LazyValue):
            self.in_expr = False
            self.current_iota_expr = value
            value = value.evaluate()
        super().__setitem__(key, value)

class EnumType(type):
    @staticmethod
    def __prepare__(name, bases, **kwargs):
        return AutoIncrementDict()

    def __new__(cls, name, bases, d):
        d = dict(d)
        d['__slots__'] = ()
        rv = type.__new__(cls, name, bases, d)
        reverse_mapping = dict(getattr(rv, '__values__', None) or {})
        for key, value in d.items():
            if isinstance(value, int):
                val = rv(value)
                setattr(rv, key, val)
                reverse_mapping[val] = key
        rv.__values__ = reverse_mapping
        return rv

class Enum(int, metaclass=EnumType):
    __slots__ = ()

    def __str__(self):
        return '%s.%s' % (
            self.__class__.__name__,
            self.__class__.__values__[self]
        )

    def __repr__(self):
        return str(self)

#
# Usage:
#

class Colors(Enum):
    red
    green
    blue
    purple = 10
    orange

assert(Colors.red == 0)
assert(Colors.green == 1)
assert(Colors.blue == 2)
assert(Colors.purple == 10)
assert(Colors.orange == 11)

class ByteSize(Enum):
    _ = iota
    KB = 1 << (10 * iota)
    MB
    GB
    TB
    PB
    EB
    ZB
    YB

assert(ByteSize._ == 0)
assert(ByteSize.KB == 1024 ** 1)
assert(ByteSize.MB == 1024 ** 2)
assert(ByteSize.GB == 1024 ** 3)
assert(ByteSize.TB == 1024 ** 4)
assert(ByteSize.PB == 1024 ** 5)
assert(ByteSize.EB == 1024 ** 6)
assert(ByteSize.ZB == 1024 ** 7)
assert(ByteSize.YB == 1024 ** 8)

class Numbers(Enum):
    zero
    one = iota
    two = iota
    six = iota * 2
    eight
    ten
    eleven = iota * 2 - 1
    thirteen
    seventy_two = iota * iota
    one_hundred_and_ten

assert(Numbers.zero == 0)
assert(Numbers.one == 1)
assert(Numbers.two == 2)
assert(Numbers.six == 6)
assert(Numbers.eight == 8)
assert(Numbers.ten == 10)
assert(Numbers.eleven == 11)
assert(Numbers.thirteen == 13)
assert(Numbers.seventy_two == 72)
assert(Numbers.one_hundred_and_ten == 110)