Refined the EnumField

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."""