From 6eea65f5814ad4a5c5c17b732e508eb78f24ea01 Mon Sep 17 00:00:00 2001
From: "Brandon T. Willard" <brandonwillard@users.noreply.github.com>
Date: Fri, 25 Jun 2021 01:01:36 -0500
Subject: [PATCH] Allow generator evaluation results

---
 etuples/core.py    | 27 +++++++++++++++++++++------
 tests/test_core.py |  8 ++++++++
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/etuples/core.py b/etuples/core.py
index b1174ce..d029008 100644
--- a/etuples/core.py
+++ b/etuples/core.py
@@ -9,6 +9,13 @@
 etuple_repr.maxother = 100
 
 
+class IgnoredGenerator:
+    __slots__ = ("gen",)
+
+    def __init__(self, gen):
+        self.gen = gen
+
+
 def trampoline_eval(z, res_filter=None):
     """Evaluate a stream of generators.
 
@@ -19,6 +26,8 @@ def trampoline_eval(z, res_filter=None):
 
     if not isinstance(z, Generator):  # pragma: no cover
         return z
+    elif isinstance(z, IgnoredGenerator):
+        return z.gen
 
     stack = deque()
     z_args, z_out = None, None
@@ -41,7 +50,10 @@ def trampoline_eval(z, res_filter=None):
         except StopIteration:
             _ = stack.pop()
 
-    return z_out
+    if isinstance(z_out, IgnoredGenerator):
+        return z_out.gen
+    else:
+        return z_out
 
 
 class InvalidExpression(Exception):
@@ -149,7 +161,8 @@ def __init__(self, seq=None, **kwargs):
     @property
     def evaled_obj(self):
         """Return the evaluation of this expression tuple."""
-        return trampoline_eval(self._eval_step())
+        res = self._eval_step()
+        return trampoline_eval(res)
 
     @property
     def eval_obj(self):
@@ -201,10 +214,12 @@ def _eval_step(self):
 
                 _evaled_obj = op(*op_args.args, **op_args.kwargs)
 
-            # assert not isinstance(_evaled_obj, ExpressionTuple)
-
-            self._evaled_obj = _evaled_obj
-            yield self._evaled_obj
+            if isinstance(_evaled_obj, Generator):
+                self._evaled_obj = _evaled_obj
+                yield IgnoredGenerator(_evaled_obj)
+            else:
+                self._evaled_obj = _evaled_obj
+                yield self._evaled_obj
 
     @evaled_obj.setter
     def evaled_obj(self, obj):
diff --git a/tests/test_core.py b/tests/test_core.py
index 80a1afb..d293508 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -1,5 +1,6 @@
 import sys
 from operator import add
+from types import GeneratorType
 
 import pytest
 
@@ -118,6 +119,13 @@ def test_op(*args):
     assert e_ladd == (1, 2, 3)
 
 
+def test_etuple_generator():
+    e_gen = etuple(lambda v: (i for i in v), range(3))
+    e_gen_res = e_gen.evaled_obj
+    assert isinstance(e_gen_res, GeneratorType)
+    assert tuple(e_gen_res) == tuple(range(3))
+
+
 def test_etuple_kwargs():
     """Test keyword arguments and default argument values."""