Last active
August 29, 2015 14:04
-
-
Save QuantumGhost/6a126e3539f208e7264b to your computer and use it in GitHub Desktop.
A simple enum implementation for Python 2
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
#!/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