public
Created

SQLAlchemy custom Enum type that uses Integer-based indexing.

  • Download Gist
sa_enum_type.py
Python
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49
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)

Please sign in to comment on this gist.

Something went wrong with that request. Please try again.