Fix JSON encoding for Enums

hexlib/db.py

@@ -2,11 +2,11 @@ import base64
 import sqlite3
 import traceback
 from datetime import datetime
+from enum import Enum
 
 import psycopg2
 import umsgpack
 from psycopg2.errorcodes import UNIQUE_VIOLATION
-import json
 from pydantic import BaseModel
 
 from hexlib.env import get_redis
@@ -15,7 +15,10 @@ from hexlib.env import get_redis
 def _json_encoder(x):
     if isinstance(x, datetime):
         return x.isoformat()
-    return x
+    if isinstance(x, Enum):
+        return x.value
+
+    raise Exception(f"I don't know how to JSON encode {x} ({type(x)})")
 
 
 class VolatileState:

setup.py

@@ -2,7 +2,7 @@ from setuptools import setup
 
 setup(
     name="hexlib",
-    version="1.83",
+    version="1.84",
     description="Misc utility methods",
     author="simon987",
     author_email="me@simon987.net",

test/test_PydanticTable.py

@@ -1,5 +1,7 @@
 import os
 from datetime import datetime
+from enum import Enum
+from typing import Optional
 from unittest import TestCase
 
 from pydantic import BaseModel
@@ -8,6 +10,11 @@ from pydantic.types import List
 from hexlib.db import PersistentState
 
 
+class Status(Enum):
+    yes = "yes"
+    no = "no"
+
+
 class Point(BaseModel):
     x: int
     y: int
@@ -16,6 +23,7 @@ class Point(BaseModel):
 class Polygon(BaseModel):
     points: List[Point] = []
     created_date: datetime
+    status: Status = Status("yes")
 
 
 class TestPydanticTable(TestCase):
@@ -35,12 +43,13 @@ class TestPydanticTable(TestCase):
             points=[
                 Point(x=1, y=2),
                 Point(x=3, y=4),
-            ]
+            ],
         )
 
         s["a"]["1"] = val
 
         self.assertEqual(s["a"]["1"].points[0].x, 1)
+        self.assertEqual(s["a"]["1"].status, Status("yes"))
         self.assertEqual(s["a"]["1"].points[1].x, 3)
         self.assertEqual(s["a"]["1"].created_date.year, 2000)