Skip to content

Instantly share code, notes, and snippets.

@Saya47
Forked from b1naryth1ef/enum.py
Created July 17, 2024 10:19
Show Gist options
  • Save Saya47/30bccda7f54a62bc61d69edbd83ac5e0 to your computer and use it in GitHub Desktop.
Save Saya47/30bccda7f54a62bc61d69edbd83ac5e0 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