diff options
Diffstat (limited to '')
-rw-r--r-- | archivist/peewee_ext.py (renamed from archivist/closure.py) | 33 |
1 files changed, 33 insertions, 0 deletions
diff --git a/archivist/closure.py b/archivist/peewee_ext.py index 01fdc75..9fda66e 100644 --- a/archivist/closure.py +++ b/archivist/peewee_ext.py @@ -1,4 +1,37 @@ from playhouse.sqlite_ext import VirtualModel, VirtualIntegerField, VirtualCharField +from peewee import Field + +class EnumField(Field): + db_field = 'enum' + + def __init__(self, enum_class, *args, **kwargs): + super().__init__(*args, **kwargs) + self.enum_class = enum_class + + def _enum_value(self, value): + if isinstance(value, str): + try: + return self.enum_class[value.upper()] + except KeyError: + pass + + try: + return self.enum_class(int(value)) + except ValueError: + raise ValueError("%r is not a valid %s" % (value, self.enum_class.__name__)) + + def db_value(self, value): + if value is None: + return value + + if isinstance(value, self.enum_class): + return value.value + + # force check of enum value + return self._enum_value(value).value + + def python_value(self, value): + return value if value is None else self._enum_value(value) def ClosureTable(model_class, referencing_class = None, foreign_key=None, id_column = None): """Model factory for the transitive closure extension.""" |