Created
April 28, 2013 22:20
-
-
Save shazow/5478648 to your computer and use it in GitHub Desktop.
SQLAlchemy custom Enum type that uses Integer-based indexing.
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
from sqlalchemy import types | |
class Enum(types.TypeDecorator): | |
impl = types.Integer | |
def __init__(self, value_map, strict=True, *args, **kw): | |
"""Emulate Enum type with integer-based indexing. | |
value_map: | |
An integer_id:name dictionary of possible values, or a list of value | |
names (which gets converted to corresponding index numbers starting from 1). | |
strict: | |
Assert that data read from the database matches with the expected | |
valid value definitions. | |
""" | |
self.strict = strict | |
if isinstance(value_map, list): | |
value_map = dict((k+1,v) for k,v in enumerate(value_map)) | |
# Enum lookup indices | |
self.id_lookup = value_map | |
self.name_lookup = dict((v,k) for k,v in value_map.iteritems()) | |
super(Enum, self).__init__() | |
def process_bind_param(self, value, dialect): | |
if value is None: | |
return | |
id = self.name_lookup.get(value) | |
if not id: | |
raise AssertionError("Name '{0}' is not one of: {1}".format(value, self.name_lookup.keys())) | |
return id | |
def process_result_value(self, value, dialect): | |
if value is None: | |
return | |
name = self.id_lookup.get(value) | |
if self.strict and not name: | |
raise AssertionError("Id '{0}' is not one of: {1}".format(value, self.id_lookup.keys())) | |
return name | |
def copy_value(self, value): | |
"Convert named value to internal id representation" | |
return self.name_lookup.get(value) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment