Search code examples
python-3.xgeneratorreturn-valueyield

Python yield: return or access other values computed within a generator


I would like to keep track of certain values (say count, stats) while the generator yields something else. My current approach is to pass a mutable object to the generator. Returning those values with "return" didn't seem to work. I was curious if there were other ways to do it.

In the following example, the generator yields lines of code from the current script and the dictionary tracker also keeps track of the count:

import sys

def gen_lines_of_code(tracker):
    loc = 0
    with open(sys.argv[0]) as fp:
        for line in fp:
            if len(line.strip()) > 0:
                loc += 1
                yield line

    tracker["count"] = loc

# Dictionary to keep track of count of lines of code
track = {"count": 0}

g = gen_lines_of_code(track)

for v in g:
    print(v, end="")

print(f"\n\tLines of code in {sys.argv[0]}: {track['count']}")

Solution

  • How about yielding tuples consisting of the line and loc?

    import io
    
    def gen_lines_of_code(file):
        loc = 0
        for line in file:
            if line.strip():
                loc += 1
                yield line, loc
    
    file = io.StringIO("""
    Lorem ipsum dolor sit amet, consectetur adipiscing elit,
    sed do eiusmod tempor incididunt ut labore et dolore magna
    aliqua.
    
    Ut enim ad minim veniam, quis nostrud exercitation ullamco
    laboris nisi ut aliquip ex ea commodo consequat.
    """)
    
    g = gen_lines_of_code(file)
    
    for v, loc in g:
        print(v, end='')
    
    try:
        print("\n\tLines of code:", loc)
    except NameError:
        loc = 0
        print("\n\tLines of code:", loc)
    

    Or you could use a iterable class (with a __iter__ method):

    class GenLinesOfCode:
    
        def __init__(self, file):
            self.file = file
            self.loc = 0
    
        def __iter__(self):
            for line in self.file:
                if line.strip():
                    self.loc += 1
                    yield line
    
    
    g = GenLinesOfCode(file)
    
    for v in g:
        print(v, end='')
    
    print("\n\tLines of code:", g.loc)