Peewee 3 Enum Field Postgres

peewee_enum.py

from peewee import *

database = PostgresqlDatabase(
    'my_database',
    **{'user': 'postgres', 'password': 'postgres', 'host': 'localhost', 'port': 5432, 'autorollback': True}
)


class BaseModel(Model):
    class Meta:
        database = database

    @classmethod
    def create_table(cls, *args, **kwargs):
        for field in cls._meta.fields.values():
            if hasattr(field, "pre_field_create"):
                field.pre_field_create(cls)

        cls._schema.create_table()


class EnumField(Field):

    def pre_field_create(self, model):
        field = "e_%s" % self.name

        self.model._meta.database.execute_sql(
            "DROP TABLE IF EXISTS %s;" % self.model._meta.table_name
        )
        
        self.model._meta.database.execute_sql(
            "DROP TYPE IF EXISTS %s;" % field
        )
        
        tail = ', '.join(["'%s'"] * len(self.choices)) % tuple(self.choices)
        q = "CREATE TYPE %s AS ENUM (%s);" % (field, tail)
        self.model._meta.database.execute_sql(q)
        
        self.field_type = "e_%s" % self.name


class MyTable(BaseModel):
    type = EnumField(choices=['one', 'two', 'three'])

    class Meta:
        table_name = 'my_table'


database.connect()
database.create_tables([MyTable])
MyTable.create(type='two')
MyTable.get(id=1).type #two
database.close()