diff --git a/hexlib/db.py b/hexlib/db.py index 79470ff..cce7405 100644 --- a/hexlib/db.py +++ b/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: diff --git a/setup.py b/setup.py index e53efe1..6a1fa7f 100644 --- a/setup.py +++ b/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", diff --git a/test/test_PydanticTable.py b/test/test_PydanticTable.py index d9c86b2..4ae85bb 100644 --- a/test/test_PydanticTable.py +++ b/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)