diff options
Diffstat (limited to 'src/mlia/core/reporting.py')
-rw-r--r-- | src/mlia/core/reporting.py | 6 |
1 files changed, 5 insertions, 1 deletions
diff --git a/src/mlia/core/reporting.py b/src/mlia/core/reporting.py index 19644b2..ad63d62 100644 --- a/src/mlia/core/reporting.py +++ b/src/mlia/core/reporting.py @@ -11,6 +11,7 @@ from collections import defaultdict from contextlib import contextmanager from contextlib import ExitStack from dataclasses import dataclass +from enum import Enum from functools import partial from io import TextIOWrapper from pathlib import Path @@ -491,13 +492,16 @@ class CustomJSONEncoder(json.JSONEncoder): """Custom JSON encoder.""" def default(self, o: Any) -> Any: - """Support numpy types.""" + """Support custom types.""" if isinstance(o, np.integer): return int(o) if isinstance(o, np.floating): return float(o) + if isinstance(o, Enum) and isinstance(o.value, str): + return o.value + return json.JSONEncoder.default(self, o) |