Last active
July 21, 2022 16:32
-
-
Save erikbern/01ae78d15f89edfa7f77e5c0a827a94d to your computer and use it in GitHub Desktop.
Output a warning if a generator function is called but never consumed
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # Let's say you have a method `map(f, inputs)` that performs some operation on inputs | |
| # Let's also sayu that this function returns a generator. | |
| # A user might pass a function f that has some side effect, e.g. `def f(x): print(x)` | |
| # This might lead to confusing results if the output is never consumed: | |
| # `map(f, inputs)` -> this does nothing, can be confusing to the user | |
| # `list(map(f, inputs))` -> this executes `f` on all inputs | |
| # To remedy the situation, we can provide a helpful warning to the user if the generator | |
| # is never consumed. | |
| ### Helper code: | |
| import functools | |
| import warnings | |
| class GeneratorWrapper: | |
| def __init__(self, gen, f): | |
| self.gen = gen | |
| self.f = f | |
| self.iterated = False | |
| def __iter__(self): | |
| self.iterated = True | |
| return self.gen | |
| def __del__(self): | |
| if not self.iterated: | |
| warnings.warn(f"Function `{self.f.__name__}` returned generator that was never consumed", UserWarning) | |
| def warn_if_not_consumed(f): | |
| @functools.wraps(f) | |
| def f_wrapped(*args, **kwargs): | |
| gen = f(*args, **kwargs) | |
| return GeneratorWrapper(gen, f) | |
| return f_wrapped | |
| ### Example usage: | |
| @warn_if_not_consumed | |
| def generator(n): | |
| for i in range(n): | |
| yield i | |
| # This doesn't warn! | |
| for x in generator(3): | |
| print(x) | |
| # This will warn! | |
| generator(5) |
Yeah, I agree that this could be extended to handle partial consumption too
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
This throws as warning if
iter()isn't called on the wrapped generator, but that's not sufficient to guarantee the generator is fully consumed. It doesn't guaranteenext()is called on the returned iterator until it has thrownStopIteration.