From 122c02b641f7185fe8585ac2b33c5070ab439d91 Mon Sep 17 00:00:00 2001
From: Zhiyuan Chen <this@zyc.ai>
Date: Fri, 23 Aug 2024 03:19:29 +0800
Subject: [PATCH] add StrEnum to yaml representer

Signed-off-by: Zhiyuan Chen <this@zyc.ai>
---
 chanfig/__init__.py | 38 +++++++++++++++++++++++++-------------
 tests/test_dump.py  | 45 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 13 deletions(-)
 create mode 100644 tests/test_dump.py

diff --git a/chanfig/__init__.py b/chanfig/__init__.py
index 130f594d..47ed8654 100755
--- a/chanfig/__init__.py
+++ b/chanfig/__init__.py
@@ -17,9 +17,21 @@
 
 from __future__ import annotations
 
-from yaml import add_representer
+from yaml import add_multi_representer
 from yaml.representer import SafeRepresenter
 
+try:
+    from enum import StrEnum
+except ImportError:
+    StrEnum = None
+
+try:
+    from strenum import LowercaseStrEnum
+    from strenum import StrEnum as UppercaseStrEnum
+except ImportError:
+    UppercaseStrEnum = None
+    LowercaseStrEnum = None
+
 from . import utils
 from ._version import __version__, __version_tuple__, version
 from .config import Config
@@ -55,15 +67,15 @@
 ]
 
 
-add_representer(FlatDict, SafeRepresenter.represent_dict)
-add_representer(NestedDict, SafeRepresenter.represent_dict)
-add_representer(DefaultDict, SafeRepresenter.represent_dict)
-add_representer(Config, SafeRepresenter.represent_dict)
-add_representer(Registry, SafeRepresenter.represent_dict)
-add_representer(ConfigRegistry, SafeRepresenter.represent_dict)
-SafeRepresenter.add_representer(FlatDict, SafeRepresenter.represent_dict)
-SafeRepresenter.add_representer(NestedDict, SafeRepresenter.represent_dict)
-SafeRepresenter.add_representer(DefaultDict, SafeRepresenter.represent_dict)
-SafeRepresenter.add_representer(Config, SafeRepresenter.represent_dict)
-SafeRepresenter.add_representer(Registry, SafeRepresenter.represent_dict)
-SafeRepresenter.add_representer(ConfigRegistry, SafeRepresenter.represent_dict)
+add_multi_representer(FlatDict, SafeRepresenter.represent_dict)
+SafeRepresenter.add_multi_representer(FlatDict, SafeRepresenter.represent_dict)
+
+if StrEnum:
+    add_multi_representer(StrEnum, SafeRepresenter.represent_str)
+    SafeRepresenter.add_multi_representer(StrEnum, SafeRepresenter.represent_str)
+if UppercaseStrEnum:
+    add_multi_representer(UppercaseStrEnum, SafeRepresenter.represent_str)
+    SafeRepresenter.add_multi_representer(UppercaseStrEnum, SafeRepresenter.represent_str)
+if LowercaseStrEnum:
+    add_multi_representer(LowercaseStrEnum, SafeRepresenter.represent_str)
+    SafeRepresenter.add_multi_representer(LowercaseStrEnum, SafeRepresenter.represent_str)
diff --git a/tests/test_dump.py b/tests/test_dump.py
new file mode 100644
index 00000000..f61023d3
--- /dev/null
+++ b/tests/test_dump.py
@@ -0,0 +1,45 @@
+# CHANfiG, Easier Configuration.
+# Copyright (c) 2022-Present, CHANfiG Contributors
+
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the following licenses:
+# - The Unlicense
+# - GNU Affero General Public License v3.0 or later
+# - GNU General Public License v2.0 or later
+# - BSD 4-Clause "Original" or "Old" License
+# - MIT License
+# - Apache License 2.0
+
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+# See the LICENSE file for more details.
+
+
+from chanfig import NestedDict
+from enum import auto
+
+try:
+    from enum import StrEnum
+except ImportError:
+    from strenum import LowercaseStrEnum as StrEnum  # type: ignore[no-redef]
+
+
+class Task(StrEnum):
+    __test__ = False
+    Regression = auto()
+    Binary = auto()
+    MultiClass = auto()
+    MultiLabel = auto()
+
+
+class TaskConfig(NestedDict):
+    __test__ = False
+    task: Task = "regression"
+
+
+class TestDump:
+    def test_yamls(self):
+        config = TaskConfig()
+        s = config.yamls()
+        assert s == "task: regression\n"