CODEWITHSUNDEEP
python-3.xgeneratorreturn-valueyield
yang5
yang5

Reputation: 1235

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']}")

Upvotes: 2

Views: 282

Answers (1)

skrx
skrx

Reputation: 20488

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)

Upvotes: 1

Related Questions