erikbern · July 21, 2022 16:32 · erikbern · Jul 21, 2022
diff --git a/warn_if_generator_is_not_consumed.py b/warn_if_generator_is_not_consumed.py
 # 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)
	# 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)