From 30737281b385e86e71c2afa1445c93afc788cf39 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Bj=C3=B6rn=20Ricks?= <bjoern.ricks@greenbone.net>
Date: Thu, 16 Nov 2023 16:33:06 +0100
Subject: [PATCH] Add: Introduce a StrEnum class for Models

Allow to use a Enum class with str behavior for easier usage of models
with enum values in APIs.
---
 pontos/models/__init__.py    | 12 ++++++++++++
 tests/models/test_strenum.py | 26 ++++++++++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 tests/models/test_strenum.py

diff --git a/pontos/models/__init__.py b/pontos/models/__init__.py
index 372b7dfc..ce432e4b 100644
--- a/pontos/models/__init__.py
+++ b/pontos/models/__init__.py
@@ -17,6 +17,7 @@
 
 from dataclasses import dataclass
 from datetime import date, datetime
+from enum import Enum
 from inspect import isclass
 from typing import Any, Dict, Type, Union, get_args, get_origin, get_type_hints
 
@@ -27,6 +28,7 @@
 __all__ = (
     "Model",
     "ModelError",
+    "StrEnum",
     "dotted_attributes",
 )
 
@@ -37,6 +39,16 @@ class ModelError(PontosError):
     """
 
 
+class StrEnum(str, Enum):
+    # Should be replaced by enum.StrEnum when we require Python >= 3.11
+    """
+    An Enum that provides str like behavior
+    """
+
+    def __str__(self) -> str:
+        return self.value
+
+
 def dotted_attributes(obj: Any, data: Dict[str, Any]) -> Any:
     """
     Set dotted attributes on an object
diff --git a/tests/models/test_strenum.py b/tests/models/test_strenum.py
new file mode 100644
index 00000000..4cf823b4
--- /dev/null
+++ b/tests/models/test_strenum.py
@@ -0,0 +1,26 @@
+# SPDX-FileCopyrightText: 2023 Greenbone AG
+#
+# SPDX-License-Identifier: GPL-3.0-or-later
+
+import unittest
+
+from pontos.models import StrEnum
+
+
+class FooEnum(StrEnum):
+    A = "a"
+    B = "b"
+
+
+class StrEnumTestCase(unittest.TestCase):
+    def test_str(self):
+        self.assertEqual(str(FooEnum.A), "a")
+        self.assertEqual(str(FooEnum.B), "b")
+
+    def test_str_append(self):
+        self.assertEqual("say " + FooEnum.A, "say a")
+        self.assertEqual("say " + FooEnum.B, "say b")
+
+    def test_f_string(self):
+        self.assertEqual(f"say {FooEnum.A}", "say a")
+        self.assertEqual(f"say {FooEnum.B}", "say b")