From 536849c98ae3e75026ead822542b936e272d2b2b Mon Sep 17 00:00:00 2001
From: maybe-sybr <58414429+maybe-sybr@users.noreply.github.com>
Date: Mon, 14 Jun 2021 17:31:20 +1000
Subject: [PATCH] Ensure regen utility class gets marked as done when
 concretised (#6789)

* fix: `regen.data` property now marks self as done

Fixes: #6786

* improv: Don't concretise regen on `repr()`

This ensures that the generator remains lazy if it's passed to `repr()`,
e.g. for logging or something.

* test: Add failing test for regen duping on errors

* refac: Remove unnecessary try in `regen.data`
---
 celery/utils/functional.py      | 14 ++++++---
 t/unit/utils/test_functional.py | 56 +++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 4 deletions(-)

diff --git a/celery/utils/functional.py b/celery/utils/functional.py
index ddf4c10379d..a82991b2437 100644
--- a/celery/utils/functional.py
+++ b/celery/utils/functional.py
@@ -241,12 +241,18 @@ def __bool__(self):
 
     @property
     def data(self):
-        try:
-            self.__consumed.extend(list(self.__it))
-        except StopIteration:
-            pass
+        if not self.__done:
+            self.__consumed.extend(self.__it)
+            self.__done = True
         return self.__consumed
 
+    def __repr__(self):
+        return "<{}: [{}{}]>".format(
+            self.__class__.__name__,
+            ", ".join(repr(e) for e in self.__consumed),
+            "..." if not self.__done else "",
+        )
+
 
 def _argsfromspec(spec, replace_defaults=True):
     if spec.defaults:
diff --git a/t/unit/utils/test_functional.py b/t/unit/utils/test_functional.py
index 58ed115b694..d7e8b686f5e 100644
--- a/t/unit/utils/test_functional.py
+++ b/t/unit/utils/test_functional.py
@@ -1,3 +1,5 @@
+import collections
+
 import pytest
 from kombu.utils.functional import lazy
 
@@ -150,6 +152,60 @@ def build_generator():
     def test_nonzero__empty_iter(self):
         assert not regen(iter([]))
 
+    def test_deque(self):
+        original_list = [42]
+        d = collections.deque(original_list)
+        # Confirm that concretising a `regen()` instance repeatedly for an
+        # equality check always returns the original list
+        g = regen(d)
+        assert g == original_list
+        assert g == original_list
+
+    def test_repr(self):
+        def die():
+            raise AssertionError("Generator died")
+            yield None
+
+        # Confirm that `regen()` instances are not concretised when represented
+        g = regen(die())
+        assert "..." in repr(g)
+
+    def test_partial_reconcretisation(self):
+        class WeirdIterator():
+            def __init__(self, iter_):
+                self.iter_ = iter_
+                self._errored = False
+
+            def __iter__(self):
+                yield from self.iter_
+                if not self._errored:
+                    try:
+                        # This should stop the regen instance from marking
+                        # itself as being done
+                        raise AssertionError("Iterator errored")
+                    finally:
+                        self._errored = True
+
+        original_list = list(range(42))
+        g = regen(WeirdIterator(original_list))
+        iter_g = iter(g)
+        for e in original_list:
+            assert e == next(iter_g)
+        with pytest.raises(AssertionError, match="Iterator errored"):
+            next(iter_g)
+        # The following checks are for the known "misbehaviour"
+        assert getattr(g, "_regen__done") is False
+        # If the `regen()` instance doesn't think it's done then it'll dupe the
+        # elements from the underlying iterator if it can be re-used
+        iter_g = iter(g)
+        for e in original_list * 2:
+            assert next(iter_g) == e
+        with pytest.raises(StopIteration):
+            next(iter_g)
+        assert getattr(g, "_regen__done") is True
+        # Finally we xfail this test to keep track of it
+        raise pytest.xfail(reason="#6794")
+
 
 class test_head_from_fun: