-
-
Save flowblok/5938793 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
# 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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment