v2.26.4

CHANGELOG.md

+# v2.26.4
+
+* `run_when_iterator_completes` and `RunActionAfterGeneratorCompletes` will now support exceptions
+
 # v2.26.3
 
 * added `run_when_iterator_completes`

satella/__init__.py

-__version__ = '2.26.3'
+__version__ = '2.26.4'

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

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)

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())