Question

I have a function in python whose output is a generator :

def main_function(x):
    r = get_range()
    for i in range(r):
        yield x+i

I want to refactor the code (I’ve simplified the use case but actual computation might be complex & longer. Please see EDIT below). Based on my understanding, this is what I should do to keep the functionality unchanged :

(a) Same interface as original code

    def sub_function(x,r):
        for i in range(r):
            yield x+i    

    def main_function(x):
        r = get_range()
        return sub_function(x,r)

As compared to other approaches where :

(b) This would return a generator of a generator (Are there any advantages of this approach ?)

    def sub_function(x,r):
        for i in range(r):
            yield x+i    

    def main_function(x):
        r = get_range()
        yield sub_function(x,r)

(c) This would defeat the purpose of a generator (Is that correct?)

    def sub_function(x,r):
        return [x+i for i in range(r)]

    def main_function(x):
        r = get_range()
        for i in sub_function(x,r):
            yield(i)

EDIT : Comments point out that the right answer is use case dependent. I want to add that my use case is parsing an XML file to extract fields and write them to a database. This part is delegated to sub_function(). I also asked this question for a general understanding of the usage of nested yield for refactoring code.

Answer 1

You’re right; the initial example and a) do the same thing since both return a generator.

b) is different: It returns a generator which yields a single element (which is another generator). To use that, you need two loops (one over the outer and one over the inner generator).

There are no advantages per se but sometimes, it can be useful to build nested generators.

c) could be worse but I’m pretty sure that the [x for x in y] is actually implemented as a generator, too. So it’s a bit more expensive but not that much.