diff --git a/CHANGELOG.md b/CHANGELOG.md index 647295eca1fd615e5081d0770a7bea0254d6395e..3f59686c5edd37927e12d31f65ac6e1159eeaab2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,7 @@ +# v2.26.4 + +* `run_when_iterator_completes` and `RunActionAfterGeneratorCompletes` will now support exceptions + # v2.26.3 * added `run_when_iterator_completes` diff --git a/satella/__init__.py b/satella/__init__.py index 9702dba3cbb93430d03e1d9779d28036507ce2f7..7ee03739b226b765ed82b8311718cafa0c0fe772 100644 --- a/satella/__init__.py +++ b/satella/__init__.py @@ -1 +1 @@ -__version__ = '2.26.3' +__version__ = '2.26.4' diff --git a/satella/coding/generators.py b/satella/coding/generators.py index 3bbe95464fdb8f635ea6e6fc059cc4195e9e632b..3b834c824b35c97cc2f6f45910abe3d80adae8ce 100644 --- a/satella/coding/generators.py +++ b/satella/coding/generators.py @@ -14,11 +14,13 @@ class RunActionAfterGeneratorCompletes(tp.Generator, metaclass=ABCMeta): __slots__ = 'generator', 'args', 'kwargs', 'closed', 'call_despite_closed' - def __init__(self, generator: tp.Generator, *args, call_despite_closed: bool = False, **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 call_on_exception: callable/1 with exception instance if generator somehow fails :param kwargs: keyword arguments to invoke action_to_run with """ self.closed = False @@ -41,6 +43,8 @@ class RunActionAfterGeneratorCompletes(tp.Generator, metaclass=ABCMeta): except StopIteration: self._try_action_run() raise + except Exception as e: + self.call_on_exception(e) def next(self): return self.generator.__next__() @@ -57,6 +61,8 @@ class RunActionAfterGeneratorCompletes(tp.Generator, metaclass=ABCMeta): except StopIteration: self._try_action_run() raise + except Exception as e: + self.call_on_exception(e) def _try_action_run(self): if not self.closed and not self.call_despite_closed: @@ -66,8 +72,11 @@ class RunActionAfterGeneratorCompletes(tp.Generator, metaclass=ABCMeta): def action_to_run(self, *args, **kwargs): """This will run when this generator completes. Override it.""" + def call_on_exception(self, exc: Exception): + """This will run when this generator throws any exception. Override it.""" -def run_when_generator_completes(gen: tp.Generator, call_on_done: tp.Callable, + +def run_when_generator_completes(gen: tp.Generator, call_on_done: tp.Callable *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 1dee1c54ac9d84852d48ee9aa0905c4594f1034a..3532dc0614ac7c25e6149bdcff06a77627d67bf8 100644 --- a/satella/coding/iterators.py +++ b/satella/coding/iterators.py @@ -138,14 +138,19 @@ class hint_with_length: return self.length -def run_when_iterator_completes(iterator: tp.Iterator, func_to_run: tp.Callable, *args, **kwargs): +def run_when_iterator_completes(iterator: tp.Iterator, func_to_run: tp.Callable, do_exception=lambda e: None, + *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 do_exception: a callable to call with the exception instance if generator fails at some point :param args: arguments to pass to the function :param kwargs: keyword arguments to pass to the function """ - yield from iterator + try: + yield from iterator + except Exception as e: + do_exception(e) func_to_run(*args, **kwargs) diff --git a/tests/test_coding/test_iterators.py b/tests/test_coding/test_iterators.py index 539b5e29a9cf86ae5a619f43d3654d68f540aa14..82459e3c2108403a9e8581265be0a22c639e77cf 100644 --- a/tests/test_coding/test_iterators.py +++ b/tests/test_coding/test_iterators.py @@ -4,7 +4,7 @@ import logging import unittest from satella.coding import SelfClosingGenerator, hint_with_length, chain, run_when_generator_completes, typing, \ - run_when_iterator_completes + run_when_iterator_completes, RunActionAfterGeneratorCompletes from satella.coding.sequences import smart_enumerate, ConstruableIterator, walk, \ IteratorListAdapter, is_empty, ListWrapperIterator @@ -74,6 +74,32 @@ class TestIterators(unittest.TestCase): pass self.assertTrue(called) + def test_run_when_iterator_fails(self): + called = False + + def generator(): + yield 1 + yield 2 + raise ValueError() + yield 3 + + def mark_done(e): + assert isinstance(e, ValueError) + nonlocal called + called = True + + def no_op(): + pass + + a = run_when_iterator_completes(generator(), no_op, mark_done) + self.assertFalse(called) + next(a) + self.assertFalse(called) + for i in a: + pass + self.assertTrue(called) + + def test_run_when_iterator_completes(self): called = False @@ -113,6 +139,33 @@ class TestIterators(unittest.TestCase): self.assertRaises(StopIteration, next, gen) self.assertFalse(called) + + def test_run_when_generator_closed_failure(self): + called = False + + def generator(): + yield 1 + yield 2 + raise ValueError() + yield 3 + + def no_op(): + pass + + class Inner(RunActionAfterGeneratorCompletes): + def action_to_run(self, *args, **kwargs): + pass + + def call_on_exception(self, exc: Exception): + nonlocal called + called = True + + gen = Inner(generator()) + a = next(gen) + gen.close() + self.assertRaises(StopIteration, next, gen) + self.assertFalse(called) + def test_list_wrapper_iterator_contains(self): lwe = ListWrapperIterator(iterate())