Python Generators Close and Return Enhancement

In my previous post I mentioned how I find it strange that the value returned by generator.close() (if it returns a value) is not made available in the next call to generator.next() (or generator.send()) as the value of the StopIteration exception. A few months ago I wrote about the problems with for loops for getting access to the value returned by a generator, and provided a solution. Well, I've improved such solution to have generator object that wraps the original generator and addresses both issues. If we close the generator and it returns a value, the next time we call next() or send() such value is available in StopIteration.value. Additionally, if a generator returns a value (either if it's closed or it finishes normally) that return value is made accessible in a result attribute of our Generator wrapper. OK, much talk, show me the code:

import inspect
from typing import Generator, Any

def cities_gen_fn():
    try:
        yield "Xixon"
        yield "Paris"
        yield "Lisbon"
        yield "Bilbao"
    except GeneratorExit:
        pass    
    # return this value both if closed or in a normal execution
    return "XPLB"


# wraps a generator in an "extended generator" that stores the value returned by close() to "return it" it the next call to .next() or .send() 
# it also stores the returned value if the original generator returns something
# that stored return-close value is only returned as StopIteration.value in the first call to .next()-.send(), ensuing calls return StopIteration.value as None

class GeneratorWithEnhancedReturnAndClose:
    def __init__(self, generator_ob: Generator[Any, Any, Any]):
        self.generator_ob = generator_ob
        self.result = None
        self.just_closed = False
        self.closed = False
        
    def __iter__(self):
        return self
    
    def _do_iterate(self, caller: str, val: Any) -> Any:
        if self.just_closed:
            self.just_closed = False
            ex = StopIteration()
            ex.value = self.result
            raise ex
        try:
            if caller == "__next__":
                return next(self.generator_ob)
            else:
                return self.generator_ob.send(val)
        except StopIteration as ex:
            if self.result is None:
                self.result = ex.value
            raise ex
            
    def __next__(self):
        return self._do_iterate("__next__", None)

    def send(self, val):
        return self._do_iterate("send", val)       

    def close(self):
        if not self.closed:
            self.closed = True
            self.just_closed = True 
            self.result = self.generator_ob.close()
            return self.result
        
    def throw(self, ex):
        return self.generator_ob.throw(ex)

print("- getting return value after for-loop")
cities = GeneratorWithEnhancedReturnAndClose(cities_gen_fn())
for city in cities:
    print(city)
print(f"return value: {cities.result}")

print("------------------------")
print("- using next() and close()")

cities = GeneratorWithEnhancedReturnAndClose(cities_gen_fn())
print(next(cities))
print(next(cities))
print(f"closing generator: {cities.close()}")
# first iteration after closing it returns the close-value in the StopIteration.value
try:
    next(cities)
except Exception as ex:
    print(f"generator finished {ex.value}")

# next iteration returns StopIteration with value = None
try:
    next(cities)
except Exception as ex:
    print(f"generator finished {ex.value}")

print(f"return value: {cities.result}")

print("------------------------")
print("- using send() and close()")
# test now that send() also works OK

def freak_cities_gen():
    try:
        w = yield "Xixon"
        w = yield f"{w}Paris{w}"
        w = yield f"{w}Lisbon{w}"
        yield f"{w}Bilbao{w}"
    except BaseException: #GeneratorExit:
        pass    
    # return this value both if closed or in a normal execution
    return "XPLB"
 
cities = GeneratorWithEnhancedReturnAndClose(freak_cities_gen())
print(next(cities))
print(cities.send("||"))
print(f"closing generator: {cities.close()}")
# first iteration after closing it returns the close-value in the StopIteration.value
try:
    next(cities) #it's the same using next or send
except Exception as ex:
    print(f"generator finished {ex.value}")

# next iteration returns StopIteration with value = None
try:
    cities.send("|") #it's the same using next or send
except Exception as ex:
    print(f"generator finished {ex.value}")

print(f"return value: {cities.result}")



# - getting return value after for-loop
# Xixon
# Paris
# Lisbon
# Bilbao
# return value: XPLB
# ------------------------
# - using next() and close()
# Xixon
# Paris
# closing generator: XPLB
# generator finished XPLB
# generator finished None
# return value: XPLB
# ------------------------
# - using send() and close()
# Xixon
# ||Paris||
# closing generator: XPLB
# generator finished XPLB
# generator finished None
# return value: XPLB

Notice that we could inspect the call stack to get from what method we are being called and rewrite the code this way:

    def _do_iterate_slow(self, val: Any) -> Any:
        # this is very cool with the use of introspection to check the caller name, but that's pretty slow
        if self.just_closed:
            self.just_closed = False
            ex = StopIteration()
            ex.value = self.result
            raise ex
        try:
            if inspect.stack()[1][3] == "__next__":
                return next(self.generator_ob)
            else:
                return self.generator_ob.send(val)
        except StopIteration as ex:
            if self.result is None:
                self.result = ex.value
            raise ex
           
    def __next__(self):
	return self._do_iterate_slow(None)

    def send(self, val):
        return self._do_iterate_slow(val)

But other than being very cool, that stack access is rather slow, so we better avoid such technique.