Merge pull request #18 from mreiche/bugfix/enum-type

Fix mapping enum types

Author: mreiche

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"

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: