from app import db


class IntEnumColumn(db.TypeDecorator):
    impl = db.Integer

    def __init__(self, enum_type, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.enum_type = enum_type

    def process_bind_param(self, value, dialect):
        if isinstance(value, self.enum_type) and isinstance(value.value, int):
            return value.value
        elif isinstance(value, int):
            return self.enum_type(value).value
        elif isinstance(value, str):
            return self.enum_type[value].value
        else:
            return TypeError()

    def process_result_value(self, value, dialect):
        return self.enum_type(value)