Skip to content

Instantly share code, notes, and snippets.

@erikbern
Last active July 21, 2022 16:32
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
Star You must be signed in to star a gist
Save erikbern/01ae78d15f89edfa7f77e5c0a827a94d to your computer and use it in GitHub Desktop.
Output a warning if a generator function is called but never consumed
# 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)
@juestr
Copy link

juestr commented Jul 21, 2022

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 guarantee next() is called on the returned iterator until it has thrown StopIteration.

@erikbern
Copy link
Author

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