diff --git a/satella/coding/generators.py b/satella/coding/generators.py index ac6673c609000a27430ff3aa451d5964fd14883d..3bbe95464fdb8f635ea6e6fc059cc4195e9e632b 100644 --- a/satella/coding/generators.py +++ b/satella/coding/generators.py @@ -67,7 +67,7 @@ class RunActionAfterGeneratorCompletes(tp.Generator, metaclass=ABCMeta): """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, *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 99b1b2427f9daa805ccf9cdad52c822d6ba958a2..1dee1c54ac9d84852d48ee9aa0905c4594f1034a 100644 --- a/satella/coding/iterators.py +++ b/satella/coding/iterators.py @@ -138,7 +138,7 @@ class hint_with_length: return self.length -def run_when_iterator_completes(iterator, func_to_run, *args, **kwargs): +def run_when_iterator_completes(iterator: tp.Iterator, func_to_run: tp.Callable, *args, **kwargs): """ Schedule a function to be called when an iterator completes.