From 8381650b04bf82a74d6fa6a8a1a9c29280e5fd40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Ko=CC=88hler?= <info@martinjankoehler.com>
Date: Mon, 16 Dec 2024 19:55:36 +0100
Subject: [PATCH] argparse helpers: for StrEnums, use the member value instead
 of the member name

---
 klayout_pex/util/argparse_helpers.py | 23 +++++++--
 tests/util/argparse_helpers_test.py  | 77 ++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 4 deletions(-)
 create mode 100644 tests/util/argparse_helpers_test.py

diff --git a/klayout_pex/util/argparse_helpers.py b/klayout_pex/util/argparse_helpers.py
index 9535da98..d945e138 100644
--- a/klayout_pex/util/argparse_helpers.py
+++ b/klayout_pex/util/argparse_helpers.py
@@ -24,18 +24,33 @@
 #
 
 import argparse
-from enum import Enum
+from enum import Enum, StrEnum
 from typing import *
 
 
 def render_enum_help(topic: str,
                      enum_cls: Type[Enum],
-                     print_default: bool = True) -> str:
+                     print_default: bool = True,
+                     lowercase_strenum: bool = False) -> str:
+    def canonic_string(name: str, member: str) -> str:
+        if issubclass(enum_cls, StrEnum):
+            if name.lower() == 'default':
+                return 'default'
+            return member.lower() if lowercase_strenum else member
+        return name.lower()
     if not hasattr(enum_cls, 'DEFAULT'):
         print_default = False
-    enum_help = f"{topic} ∈ {set([name.lower() for name, member in enum_cls.__members__.items()])}"
+    case_list = [f"'{canonic_string(name, member)}'" for name, member in enum_cls.__members__.items()]
+    enum_help = f"{topic} ∈ \u007b{', '.join(case_list)}\u007d"
     if print_default:
-        enum_help += f".\nDefaults to '{getattr(enum_cls, 'DEFAULT').name.lower()}'"
+        default_case: enum_cls = getattr(enum_cls, 'DEFAULT')
+        if issubclass(enum_cls, StrEnum):
+            default_value: str = default_case.value
+            if lowercase_strenum:
+                default_value = default_value.lower()
+        else:
+            default_value = default_case.name.lower()
+        enum_help += f".\nDefaults to '{default_value}'"
     return enum_help
 
 
diff --git a/tests/util/argparse_helpers_test.py b/tests/util/argparse_helpers_test.py
new file mode 100644
index 00000000..68b0e782
--- /dev/null
+++ b/tests/util/argparse_helpers_test.py
@@ -0,0 +1,77 @@
+#
+# --------------------------------------------------------------------------------
+# SPDX-FileCopyrightText: 2024 Martin Jan Köhler and Harald Pretl
+# Johannes Kepler University, Institute for Integrated Circuits.
+#
+# This file is part of KPEX 
+# (see https://github.com/martinjankoehler/klayout-pex).
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# 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
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+# SPDX-License-Identifier: GPL-3.0-or-later
+# --------------------------------------------------------------------------------
+#
+
+from __future__ import annotations
+
+import allure
+from enum import StrEnum, IntEnum
+from typing import *
+import unittest
+
+from klayout_pex.util.argparse_helpers import render_enum_help
+
+
+@allure.parent_suite("Unit Tests")
+@allure.tag("Enum", "Help", "Util")
+class Test(unittest.TestCase):
+    def test_render_enum_help__nodefault__intenum(self):
+        class IntEnum1(IntEnum):
+            CASE1 = 1
+            CASE2 = 2
+
+        expected_string = "--arg ∈ {'case1', 'case2'}"
+        obtained_string = render_enum_help(enum_cls=IntEnum1, topic="--arg", print_default=False)
+        assert obtained_string == expected_string
+
+    def test_render_enum_help__default__intenum(self):
+        class IntEnum2(IntEnum):
+            CASE1 = 1
+            CASE2 = 2
+            DEFAULT = CASE1
+
+        expected_string = "--arg ∈ {'case1', 'case2', 'default'}.\nDefaults to 'case1'"
+        obtained_string = render_enum_help(enum_cls=IntEnum2, topic="--arg", print_default=True)
+        assert obtained_string == expected_string
+
+    def test_render_enum_help__default__strenum(self):
+        class StrEnum1(StrEnum):
+            CASE1 = "Case1"
+            CASE2 = "Case2"
+            DEFAULT = CASE1
+
+        expected_string = "--arg ∈ {'Case1', 'Case2', 'default'}.\nDefaults to 'Case1'"
+        obtained_string = render_enum_help(enum_cls=StrEnum1, topic="--arg", print_default=True)
+        assert obtained_string == expected_string
+
+    def test_render_enum_help__default__strenum__lowercased(self):
+        class StrEnum1(StrEnum):
+            CASE1 = "Case1"
+            CASE2 = "Case2"
+            DEFAULT = CASE1
+
+        expected_string = "--arg ∈ {'case1', 'case2', 'default'}.\nDefaults to 'case1'"
+        obtained_string = render_enum_help(enum_cls=StrEnum1, topic="--arg",
+                                           print_default=True,
+                                           lowercase_strenum=True)
+        assert obtained_string == expected_string