From d1b2374cda6811a93d1174400fc2eecd7100a8c3 Mon Sep 17 00:00:00 2001 From: Dmitrii Agibov Date: Mon, 30 Jan 2023 14:42:24 +0000 Subject: MLIA-785 Enable export into json for enums Change-Id: I8e4d5d04f6b1b252dae872ea76d2bd8c41f4b376 --- src/mlia/core/reporting.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) (limited to 'src/mlia/core') 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) -- cgit v1.2.1