QuantumGhost
7/25/2014 - 3:31 PM

A simple enum implementation for Python 2

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))