From 9caec1e718513ea6c420b2daca73c2ef91e743bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Ma=C5=9Blanka?= <piotr.maslanka@henrietta.com.pl>
Date: Sun, 17 Nov 2024 20:17:57 +0100
Subject: [PATCH] added `run_when_iterator_completes`
---
CHANGELOG.md | 4 ++++
docs/coding/sequences.rst | 4 +++-
satella/__init__.py | 2 +-
satella/coding/__init__.py | 4 ++--
satella/coding/generators.py | 26 +++++++++++++++++++-------
satella/coding/iterators.py | 13 +++++++++++++
tests/test_coding/test_iterators.py | 23 ++++++++++++++++++++++-
7 files changed, 64 insertions(+), 12 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b4a0d716..647295ec 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,7 @@
+# v2.26.3
+
+* added `run_when_iterator_completes`
+
# v2.26.2
* RunActionAfterGeneratorCompletes won't call it's on_done action if closed prematurely
diff --git a/docs/coding/sequences.rst b/docs/coding/sequences.rst
index ab964020..857c5fbb 100644
--- a/docs/coding/sequences.rst
+++ b/docs/coding/sequences.rst
@@ -7,7 +7,9 @@ Generators
.. autoclass:: satella.coding.RunActionAfterGeneratorCompletes
:members:
-.. autoclass:: satella.coding.
+.. autofunction:: satella.coding.run_when_generator_completes
+
+.. autofunction:: satella.coding.run_when_iterator_completes
Rolling averages
================
diff --git a/satella/__init__.py b/satella/__init__.py
index b9d4f6d2..dc3b7a8e 100644
--- a/satella/__init__.py
+++ b/satella/__init__.py
@@ -1 +1 @@
-__version__ = '2.26.2'
+__version__ = '2.26.3a1'
diff --git a/satella/coding/__init__.py b/satella/coding/__init__.py
index 3c8acc60..653a3bb7 100644
--- a/satella/coding/__init__.py
+++ b/satella/coding/__init__.py
@@ -13,7 +13,7 @@ from .deleters import ListDeleter, DictDeleter
from .environment import Context
from .expect_exception import expect_exception
from .fun_static import static_var
-from .iterators import hint_with_length, SelfClosingGenerator, exhaust, chain
+from .iterators import hint_with_length, SelfClosingGenerator, exhaust, chain, run_when_iterator_completes
from .metaclasses import metaclass_maker, wrap_with, dont_wrap, wrap_property, DocsFromParent, \
CopyDocsFrom
from .misc import update_if_not_none, update_key_if_none, update_attr_if_none, queue_iterator, \
@@ -26,7 +26,7 @@ from .recast_exceptions import rethrow_as, silence_excs, catch_exception, log_ex
from .generators import RunActionAfterGeneratorCompletes, run_when_generator_completes
__all__ = [
- 'RunActionAfterGeneratorCompletes', 'run_when_generator_completes',
+ 'RunActionAfterGeneratorCompletes', 'run_when_generator_completes', 'run_when_iterator_completes',
'EmptyContextManager', 'Context', 'length',
'assert_equal', 'InequalityReason', 'Inequal', 'wrap_callable_in_context_manager',
'Closeable', 'contains', 'enum_value',
diff --git a/satella/coding/generators.py b/satella/coding/generators.py
index fd29de6a..ac6673c6 100644
--- a/satella/coding/generators.py
+++ b/satella/coding/generators.py
@@ -12,26 +12,35 @@ class RunActionAfterGeneratorCompletes(tp.Generator, metaclass=ABCMeta):
via close()
"""
- __slots__ = 'generator', 'args', 'kwargs', 'closed'
+ __slots__ = 'generator', 'args', 'kwargs', 'closed', 'call_despite_closed'
- def __init__(self, generator: tp.Generator, *args, **kwargs):
+ def __init__(self, generator: tp.Generator, *args, call_despite_closed: bool = False, **kwargs):
"""
:param generator: generator to watch for
:param args: arguments to invoke action_to_run with
+ :param call_despite_closed: :meth:`action_to_run` will be called even if the generator is closed
:param kwargs: keyword arguments to invoke action_to_run with
"""
self.closed = False
self.generator = generator
self.args = args
+ self.call_despite_closed = call_despite_closed
self.kwargs = kwargs
def close(self):
+ """
+ Close this generator. Note that this will cause :meth:`action_to_run` not to run
+ """
self.closed = True
self.generator.close()
def send(self, value):
"""Send a value to the generator"""
- return self.generator.send(value)
+ try:
+ return self.generator.send(value)
+ except StopIteration:
+ self._try_action_run()
+ raise
def next(self):
return self.generator.__next__()
@@ -46,16 +55,19 @@ class RunActionAfterGeneratorCompletes(tp.Generator, metaclass=ABCMeta):
try:
return self.generator.__next__()
except StopIteration:
- if not self.closed:
- self.action_to_run(*self.args, **self.kwargs)
+ self._try_action_run()
raise
+ def _try_action_run(self):
+ if not self.closed and not self.call_despite_closed:
+ self.action_to_run(*self.args, **self.kwargs)
+
@abstractmethod
- def action_to_run(self):
+ def action_to_run(self, *args, **kwargs):
"""This will run when this generator completes. Override it."""
-def run_when_generator_completes(gen: tp.Generator, call_on_done: tp.Callable[[], None],
+def run_when_generator_completes(gen: tp.Generator, call_on_done: tp.Callable[[...], None],
*args, **kwargs) -> RunActionAfterGeneratorCompletes:
"""
Return the generator with call_on_done to be called on when it finishes
diff --git a/satella/coding/iterators.py b/satella/coding/iterators.py
index 2fa23658..99b1b242 100644
--- a/satella/coding/iterators.py
+++ b/satella/coding/iterators.py
@@ -136,3 +136,16 @@ class hint_with_length:
return self.length_factory()
else:
return self.length
+
+
+def run_when_iterator_completes(iterator, func_to_run, *args, **kwargs):
+ """
+ Schedule a function to be called when an iterator completes.
+
+ :param iterator: iterator to use
+ :param func_to_run: function to run afterwards
+ :param args: arguments to pass to the function
+ :param kwargs: keyword arguments to pass to the function
+ """
+ yield from iterator
+ func_to_run(*args, **kwargs)
diff --git a/tests/test_coding/test_iterators.py b/tests/test_coding/test_iterators.py
index b1afe729..539b5e29 100644
--- a/tests/test_coding/test_iterators.py
+++ b/tests/test_coding/test_iterators.py
@@ -3,7 +3,8 @@ import sys
import logging
import unittest
-from satella.coding import SelfClosingGenerator, hint_with_length, chain, run_when_generator_completes, typing
+from satella.coding import SelfClosingGenerator, hint_with_length, chain, run_when_generator_completes, typing, \
+ run_when_iterator_completes
from satella.coding.sequences import smart_enumerate, ConstruableIterator, walk, \
IteratorListAdapter, is_empty, ListWrapperIterator
@@ -73,6 +74,26 @@ class TestIterators(unittest.TestCase):
pass
self.assertTrue(called)
+ def test_run_when_iterator_completes(self):
+ called = False
+
+ def generator():
+ yield 1
+ yield 2
+ yield 3
+
+ def mark_done():
+ nonlocal called
+ called = True
+
+ a = run_when_iterator_completes(generator(), mark_done)
+ self.assertFalse(called)
+ next(a)
+ self.assertFalse(called)
+ for i in a:
+ pass
+ self.assertTrue(called)
+
def test_run_when_generator_closed(self):
called = False
--
GitLab