Skip to content

Instantly share code, notes, and snippets.

@waltercool
Forked from b1naryth1ef/enum.py
Created May 17, 2018 15:13
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 waltercool/8d175e69c6394c2b42a889a217e1936f to your computer and use it in GitHub Desktop.
Save waltercool/8d175e69c6394c2b42a889a217e1936f to your computer and use it in GitHub Desktop.
PeeWee Postgres Enum Field
from peewee import *
class BModel(Model):
class Meta:
database = db
@classmethod
def create_table(cls, *args, **kwargs):
for field in cls._meta.get_fields():
if hasattr(field, "pre_field_create"):
field.pre_field_create(cls)
cls._meta.database.create_table(cls)
for field in cls._meta.get_fields():
if hasattr(field, "post_field_create"):
field.post_field_create(cls)
class EnumField(Field):
db_field = "enum"
def pre_field_create(self, model):
field = "e_%s" % self.name
self.get_database().get_conn().cursor().execute(
"DROP TYPE IF EXISTS %s;" % field
)
query = self.get_database().get_conn().cursor()
tail = ', '.join(["'%s'"] * len(self.choices)) % tuple(self.choices)
q = "CREATE TYPE %s AS ENUM (%s);" % (field, tail)
query.execute(q)
def post_field_create(self, model):
self.db_field = "e_%s" % self.name
def coerce(self, value):
if value not in self.choices:
raise Exception("Invalid Enum Value `%s`", value)
return str(value)
def get_column_type(self):
return "enum"
def __ddl_column__(self, ctype):
return SQL("e_%s" % self.name)
class Example(BModel):
enum = EnumField(choices=["a", "b", "c"])
Example.create(enum="a")
Example.get(id=1).enum # "a"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment