From 0ff61f8d60603f5207f9525ea74d1b585e714a43 Mon Sep 17 00:00:00 2001
From: Mike Reiche <mike@reiche.world>
Date: Fri, 5 Sep 2025 14:13:57 +0200
Subject: [PATCH] Fix mapping enum types

---
 tests/test_opt.py | 10 ++++++++++
 tinystream.py     | 14 ++++++++++----
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/tests/test_opt.py b/tests/test_opt.py
index b8a5df0..62f90a0 100644
--- a/tests/test_opt.py
+++ b/tests/test_opt.py
@@ -1,3 +1,5 @@
+from enum import Enum
+
 import pytest
 
 from test_streams import Node
@@ -169,3 +171,11 @@ def test_stream():
 def test_dict_map_kwargs():
     opt = Opt({"name": "First"})
     assert opt.map_kwargs(Node).get().name == "First"
+
+
+def test_map_enum():
+    class AnalysisAnalysisState(str, Enum):
+        NOT_SET = "NOT_SET"
+        RESOLVED = "RESOLVED"
+
+    assert Opt(AnalysisAnalysisState.RESOLVED).map_key("value").get() == "RESOLVED"
diff --git a/tinystream.py b/tinystream.py
index 7b552ba..cadf957 100644
--- a/tinystream.py
+++ b/tinystream.py
@@ -1,3 +1,4 @@
+import enum
 import functools
 import itertools
 from typing import Iterable, TypeVar, Callable, List, Dict, Tuple, Iterator, Generic, Type
@@ -18,22 +19,27 @@
 
 
 def _key_exists(x: any, key: Key, invert: bool = False):
+    def _hasattr():
+        if invert:
+            return not hasattr(x, key)
+        else:
+            return hasattr(x, key)
+
     if isinstance(x, (list, tuple)):
         size = len(x)
         if invert:
             return key >= size
         else:
             return key < size
+    elif isinstance(x, enum.Enum):
+        return _hasattr()
     elif isinstance(x, (dict, Iterable)):
         if invert:
             return key not in x
         else:
             return key in x
     else:
-        if invert:
-            return not hasattr(x, key)
-        else:
-            return hasattr(x, key)
+        return _hasattr()
 
 
 def _get_key_value(x: any, key: Key) -> any: