How to use yield in Python

There’s nothing scarier than a code that you don’t understand. A few years ago I got concerned every time that I spotted a yield command. It seemed to be returning something... Somehow it worked even though I never used that command before. Is it like return?

Well, they’re similar, but work differently for sure! You're about to learn how yield works and what the hell are generators in a FUN WAY.

You'll be able to run the python commands right away from this blog as we progress so you can not only understand but SEE how it works.

(It might be disabled if reading on AMP mode or on small screens though. If that's the case you might want to come back to this article later.)

You can open the Python environment by either hitting ctrl + '  or manually clicking on the bar at the bottom of this blog.

Executing the Python code can be achieved by either ctrl + Enter / cmd + Enter, or just click Run.

What does yield do in Python?

Any yield command automatically makes your code return a generator object. Generators are functions that can be paused and resumed.

Try this:

def mygenerator():
    n = 1
    yield n
    n += 2
    yield n

print(mygenerator())

And you’re going to see something like: <generator object mygenerator at 0xd91df0> as your output.

Not much impressive right? That's because we didn't use the generator object just yet. Let’s see what’s inside the generator we just returned.

Please, wrap the object with next:

def mygenerator():
    n = 1
    yield n
    n += 2
    yield n

print(next(mygenerator()))

Now you should see 1  as the result.

You're probably familiar with the return command and this would be the expected output if any function attempted to return twice: The first return interrupts the function and returns its results.

As we discussed before, this is not true for generators, further yield commands along the way are on hold.

Let's make it obvious by adding prints, and by forcing it to run fully with list:

def mygenerator():
    n = 1
    print("A short pause")
    yield n
    
    print("The pause is over!")
    n += 2
    yield n

print(list(mygenerator()))

You should see all our prints, and all yield values ( [1, 3]).

Cool right?

Unfortunately, it was not so obvious to see the "pauses" happening. Let's be a bit more verbose now by iterating over the generator, and doing calculations before the generator finishes:

def mygenerator():
    print("Gen: Yielding 1")
    yield 1

    print("Gen: Yielding 2")
    yield 2
    
    print("Gen: Yielding 3")
    yield 3


gen = mygenerator()
for n in gen:
    print(f"For: I got {n}")
    print("For: Since generator is paused, I can do some calculation")
    print(f"For: {n} * 10 = {n * 10}") 
    print()

Note how we can calculate (multiply by ten) even before the generator finishes processing!

My proposal here is to make things stupidly obvious. So here's how it would differ from a regular function, try this and check the output:

def myregularfunc():
    ret = []
    for n in range(1, 4):
    	print(f"Fun: Preparing {n}")
    	ret.append(n)

    print("Returning full list")
    return ret
    
ret = myregularfunc()
print()
for n in ret:
    print(f"For: I got {n}")
    print(f"For: {n} * 10 = {n * 10}") 
    print()

Note how inevitably we have to wait for myregularfunc to end in order to do anything we need with its results.

🏷 Python type hint for Generators

Do you want to find out how to type hint your generator functions? Easy!

from typing import Iterator, Iterable, Generator

def mygen1() -> Iterator[int]:
    """Iterator is fine"""
    yield 1
    
def mygen2() -> Iterable[int]:
    """Iterable is equally fine"""
    yield 2

def mygenverbose() -> Generator[str, None, None]:
    """
    Verbose way of defining a generator, it receives 3 types:
    (1) Yield type
    (2) Send type (if set, otherwise None)
    (3) Return type (if returning at the end, otherwise None)    
    """
    yield "That's verbose"

print(mygen1)
print(mygen2)
print(mygenverbose)

I'm personally used to Generator[type, None, None] for keeping things explicit. Feel free to pick whatever you prefer.

🔂 Generators of generators

There're also moments where we want to generate from generators. Using just yield won't work, see:

from typing import Generator

def starting_five() -> Generator[int, None, None]:
    """Generator that returns integers from 1-5"""
    for n in range(1, 6):
        yield n
        
        
def ending_five() -> Generator[int, None, None]:
    """Generator that returns integers from 6-10"""
    for n in range(6, 11):
        yield n
        
        
def all_ten() -> Generator[int, None, None]:
    """Generator that relies on other generators"""
    yield starting_five() # This is broken
    yield ending_five()  # This is broken
    
    
print(list(all_ten()))

By doing this you're creating a generator that returns generator objects (The type hint would be something like: Generator[Generator[int, None, None], None, None]) which is not what we want!

To get it working you have to use yield from!

Fix it yourself! Replace the bare yield ... from all_ten with yield from ...

And you should have the expected result: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

⏱ Python Generators Use Cases

I enjoy bringing to you real examples that I actually use, and not just a bunch of syntax. So here're 3 times where I had and still have to deal with generators:

1️⃣  To make your code cleaner

Have you realized that functions that build lists to return need more lines? You need to set up the list, append, and then return it. The Effective Python book recommends generators whenever you need to return a list or iterable.

from typing import List, Generator


def regularfunc() -> List[int]:
    """This works and it's fine."""
    ret = []
    for n in range(10):
        ret.append(n * 10)
        
    return ret
    

def smallergen() -> Generator[int, None, None]:
    """This also works. Favor this one."""
    for n in range(10):
        yield n * 10
        
        
# Intended Usage:
print("Regular func:", regularfunc())
print("Smaller gen:", list(smallergen()))

2️⃣  To read large files (and save memory)

Once again: Generators are functions that can be paused and resumed.

It means that when you need to load a big file like a CSV you can read it line by line, instead of loading the whole thing and wasting all your memory.

Check:

def create_dummy_file():
    with open("loadme.csv", "w") as f:
        f.write("Num,Name\n")
    
        for n in range(10):
            f.write(f"{n},John Doe\n")
            
            
def read_the_whole_file():
    with open("loadme.csv", "r") as f:
        return f.readlines()


def read_line_by_line():
    with open("loadme.csv", "r") as f: 
        for row in f:
            yield row
            

create_dummy_file()

# Bad idea when file is big:
csvcontent = read_the_whole_file()
for row in csvcontent:
    # Even though we only use one line, it already loaded all lines
    print("Eager load:", row)

# Good idea when file is big:
csvcontent = read_line_by_line()
for row in csvcontent:
    # only loads a row to process
    print("Lazy load:", row)

Even though the output is the same, whenever you're reading a file (that might be big) and requires processing, stick to generators. It may save you from some MemoryErrors.

3️⃣  When writing tests with Pytest

Pytest has the concept of fixtures (out of the scope from this article so I'll be short), and there are commands that you want to be executed before AND after a test. See:

FILECONTENT = """
import pytest

@pytest.fixture
def notifyonlybeforefixture():
    print('Before the test ends')
    
    
@pytest.fixture
def notifycompletefixture():
    print('Before the test ends')
    yield # trick!
    print('After the test ends')
    


def test_fixtures_before(notifyonlybeforefixture):
    print("During test")
    assert False
    
    
def test_fixtures_complete(notifycompletefixture):
    print("During test")
    assert False
"""

# Note: This is only needed so you can test from your browser.
with open("mytest.py", "w") as f:
    f.write(FILECONTENT)


import pytest
pytest.main(["-v", "--cache-clear", "mytest.py"])

Since you're running Python from your browser, there's no default file. We have to generate one named mytest.py , so we can execute pytest on it. I hope this quirck won't make it too hard for you to understand.

Note how just adding a yield inside our notifycompletefixture function is enough to let pytest know it's ok to execute the test, and once it's completed, get back and resume from where it was left.

That's valuable when you need to tear down things (maybe close a file, clear database state, or similar).

By the way, do you know what else can be paused and resumed? Python Async. If you don't know what I'm talking about you should totally read this article.

Thank you for reading this far! I deeply appreciate your time and I hope you had fun playing with Python without having to leave the browser!

That was my first attempt to make Python learning fun and less boring. It would be an honor for me to hear your feedback: How was it?

Hi there! Are you enjoying what you're reading? If you do, you can be the first to know when I post cool stuff!

Sure! Let me know No, thanks