DEV Community

loading...
Cover image for What does yield do in Python?

What does yield do in Python?

hmajid2301 profile image Haseeb Majid Originally published at haseebmajid.dev ・4 min read

In this article, we will go over what the yield keyword is used for. We will also cover how you can use a yield with a pytest fixture to allow us to clean up data after our tests.

Background

Iterables & Iterators

Before we can look at the yield keyword we will need to cover iterables and generators in Python. An "iterable" is any Python object that can return its members one at a time, in a for-loop.

In Python we have functions called magic methods, there are methods like __enter__ and __exit__ defined within objects. These are called "magic" methods because they are never directly called by the user. For an object to be iterable, it needs to implement the __iter__ method. If an object is iterable it can be passed to the iter() function. The iter() function returns an iterator.

 ipython
Python 3.8.5 (default, Jul 28 2020, 12:59:40)
Type 'copyright', 'credits' or 'license' for more information
IPython 7.14.0 -- An enhanced Interactive Python. Type '?' for help.

In [1]: iter([1, 2, 3])
Out[1]: <list_iterator at 0x7f4c11556730>

In [2]: iter("hello")
Out[2]: <str_iterator at 0x7f4c11598c10>

In [3]: iter(42)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-3-ef50b48e4398> in <module>
----> 1 iter(42)

TypeError: 'int' object is not iterable

In [4]:
Enter fullscreen mode Exit fullscreen mode

An iterator is any object which has the __next__ magic method defined. Whenever we use a for-loop (or list comprehension), the next method is called automatically for us, to get the next item from the iterable.

In [5]: hello_list = ["h", "e", "l", "l", "o"]

In [6]: iterator = iter(hello_list)

In [7]: next(iterator)
Out[7]: 'h'

In [8]: next(iterator)
Out[8]: 'e'

In [9]: next(iterator)
Out[9]: 'l'

In [10]: next(iterator)
Out[10]: 'l'

In [11]: next(iterator)
Out[11]: 'o'

In [12]: next(iterator)
---------------------------------------------------------------------------
StopIteration                             Traceback (most recent call last)
<ipython-input-12-4ce711c44abc> in <module>
----> 1 next(iterator)

StopIteration:
Enter fullscreen mode Exit fullscreen mode

In summary, an iterable is an object that can be "looped" over and an iterator is an object which can do the "looping" for us, it will keep track of the current state/index and move to the next item. In the example above the hello_list is iterable and the iterator variable is the iterator.

Generators

Generators are a special type of iterable, they differ from normal lists in two main ways:

  • You can only iterate over them once
  • They don't store all of their values in memory

So generators can be great when lists get very large.

In [14]: g = (x^2 for x in range(10))

In [15]: for i in g:
    ...:     print(i)
    ...:
2
3
0
1
6
7
4
5
10
11

In [16]: for i in g:
    ...:     print(i)
    ...:
Enter fullscreen mode Exit fullscreen mode

Yield

Now that we finally understand iterables and generators let's see how they relate to the yield keyword. yield can be used like return except it will return a generator.

In [17]: def example():
    ...:     yield "A"
    ...:     yield "B"
    ...:     yield "C"
    ...:

In [18]: for i in example():
    ...:     print(i)
    ...:
A
B
C

In [22]: example()
Out[22]: <generator object example at 0x7f4c1147a0b0>
Enter fullscreen mode Exit fullscreen mode

A good example of yield can be seen above, it differs from a return because it is smart enough to retain "state" and resume where it left off in the function. We can see the same example with return. In this example there is only a single item being returned so only "A" is being looped over.

In [19]: def example():
    ...:     return "A"
    ...:     return "B"
    ...:     return "C"
    ...:
    ...:

In [20]: for i in example():
    ...:     print(i)
    ...:
A
Enter fullscreen mode Exit fullscreen mode

Pytest Example

One interesting use case of using the yield keyword is using it to run clean up tasks after running tests using pytest. Pytest is a very popular testing framework in Python, it allows us to create a file called conftest.py. Here we store common functions, fixtures shared between our tests.

In the example below, before any tests have run the clean_up fixture will be called, because we have given it the autouse=True parameter. It will yield, and return a generator after all of our tests have finished running. The print and the teardown tasks will then be run. This is useful for example when you want to clean up your database after running tests that will add "test" data to it. Or in fact, any other type of teardown tasks you need to run after all of your tests have finished running.

@pytest.fixture(scope="session", autouse=True)
def clean_up():
    yield
    print("teardown after yield")
    delete_database_collection()
Enter fullscreen mode Exit fullscreen mode

Appendix

Discussion (0)

pic
Editor guide