Skip to content

Instantly share code, notes, and snippets.

@QuantumGhost
Last active August 29, 2015 14:04
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save QuantumGhost/6a126e3539f208e7264b to your computer and use it in GitHub Desktop.
Save QuantumGhost/6a126e3539f208e7264b to your computer and use it in GitHub Desktop.
A simple enum implementation for Python 2
#!/usr/bin/env python3
# coding=utf-8
from __future__ import print_function
from weakref import proxy
integer_type = (int, long)
class EnumAttr(object):
def __init__(self, cls, name, value):
# Avoid adding ref-count to the Enum class
self.cls = proxy(cls)
self.name, self.value = name, value
def __eq__(self, other):
if isinstance(other, EnumAttr):
return self.cls == other.cls and self.value == other.value
else:
return self.value == other
def __repr__(self):
return "<{}.{}: {}>".format(self.cls.__name__, self.name, self.value)
class ReadOnly(object):
"""
Read only property descriptor, used for enum class.
"""
def __init__(self, enum_attr):
self.enum_attr = enum_attr
def __get__(self, cls, meta):
return self.enum_attr
def __set__(self, cls, value):
raise AttributeError("Cannot set attribute {} for enum type {}.".format(self.enum_attr.name, cls.__name__))
def __delete__(self, cls):
raise AttributeError("Cannot delete attribute {} for enum type {}.".format(self.enum_attr.name, cls.__name__))
class EnumMeta(type):
# TODO: add unique constraint
def __init__(cls, name, bases, attrs):
cls.__reverse_map = {}
print(attrs)
for key, value in attrs.iteritems():
# Do not use descriptors with buildin attributes
if not (key.startswith('__') and key.endswith('__')):
if value not in cls.__reverse_map:
enum_attr = EnumAttr(cls, key, value)
descriptor = ReadOnly(enum_attr)
setattr(cls, key, descriptor)
cls.__reverse_map[value] = descriptor
else:
# make an alias for existed attribute
setattr(cls, key, cls.__reverse_map[value])
type.__init__(cls, name, bases, attrs)
def __call__(cls, value):
if value in cls.__reverse_map:
descriptor = cls.__reverse_map[value]
return descriptor.__get__(cls, None)
else:
raise ValueError("{} is not a valid {}.".format(value, cls.__name__))
class Enum(object):
# TODO: figure out how to handle duplicated keys.
"""
This is a base class of enum type.
"""
__metaclass__ = EnumMeta
class Fruit(Enum):
orange = 1
apple = 1
class ExtendedFruit(Fruit):
banana = 3
peach = 4
class Computer(Enum):
notebook = 1
pc = 2
if __name__ == '__main__':
print('====== Fruit type ======')
print(repr(Fruit.orange))
print(Fruit.apple)
print(Fruit(1))
print(Fruit.apple == 1)
print("====== Computer type ======")
print(repr(Computer.notebook))
print(Computer.pc)
print(Computer.pc == 1)
print(Computer.pc == 2)
print(Computer(2))
print(Computer.notebook == Fruit.apple)
print("====== Extended Fruit type ======")
print(ExtendedFruit.orange == 1)
print(ExtendedFruit.orange == Fruit.orange)
print(ExtendedFruit(3))
print(ExtendedFruit(5))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment