How to Retrieve a Single Element from a Python Generator

by
5/5 - (21 votes)

This article will show you how to retrieve a single generator element in Python. Before moving forward, let’s review what a generator does.

Quick Recap Generators

💡 Definition Generator: A generator is commonly used when processing vast amounts of data. This function uses minimal memory and produces results in less time than a standard function.

💡 Definition yield: The yield keyword returns a generator object instead of a value.

💡 Definition next(): This function takes an iterator and an optional default value. Each time this function is called, it returns the next iterator item until all items are exhausted. Once exhausted, it returns the default value passed or a StopIteration Error.

Problem Formulation and Solution Overview

💬 Question: How would we write code to retrieve and return a single element from a Generator?

We can accomplish this task by one of the following options:

Method 1: Use a Generator and random.randint()

This example uses a Generator and the random.randint() function to return a random single element.

Let’s say we want to start a weekly in-house lottery called LottoOne (yeah — the lottery naming doesn’t care about Python’s naming conventions 🐍).

The code below will generate and return one (1) random element (an integer): the winning number for the week.

import random

def LottoOne():
    num = random.randint(1, 50)
    yield print(f'The Winning Number for the Week is: {num}!')

gen = LottoOne()
next(gen)

The first line in the above code imports the random library. This library allows the generation of random numbers using the random.randint() function.

On the following line, an instance of LottoOne is instantiated and saved to the variable gen. If output to the terminal, an object similar to that shown below will display.

<generator object LottoOne at 0x00000236B9686880>

Since we only want the first randomly generated number, the code calls the next() function once and passes it one (1) argument, the object gen. The results are output to the terminal.

The Winning Number for the Week is: 44
Understanding Generators In Python

👉 Recommended Tutorial: An Introduction to Python Classes

Method 2: Use a Generator and islice()

This example uses a generator and the itertools.islice() function to return a single element.

If you know what number you need to return from a generator, it can be referenced directly, as shown in the code below.

import itertools
from itertools import islice

gen = (x for x in range(1, 50))
res = next(itertools.islice(gen, 2, None))
print(res)

The first two (2) lines in the above code import the itertools library and its associated function islice() needed to achieve the desired result.

The following line creates a generator comprehension using the range() function and passing it a start and stop position (1, 50-1). The results save to gen as an object.

If output to the terminal, an object similar to that shown below will display.

<generator object at 0x00000285D4DB68F0>

Then, the next() function is called and passed one (1) argument, itertools.islice().

This function is then passed three (3) arguments:

  1. An iterator. In this case, gen.
  2. The value to return, idx.
  3. The default value. In this case, the keyword None. Passing a default value prevents a StopIteration error from occurring when the end of the iterator has been reached.

The results save to res and are output to the terminal.

3

The value of 3 can be found at index 2 in gen.

Iterators, Iterables and Itertools

Method 3: Use a List, Generator Comprehension and slicing

This example uses a list, generator comprehension and slicing to return a single element.

gen = (i for i in range(1, 50))
res = list(gen)[3]
print(res)

The first line in the above code creates a Generator Comprehension using the range() function and passing it a start and stop position (1, 50-1). The results save to gen as an object.

If output to the terminal, an object similar to that shown below will display.

<generator object at 0x00000295D4DB78F0>

The following line converts the object to a list. Slicing is then applied to retrieve the list element at position three (3).

The results save to res and are output to the terminal.

4

The value of 4 can be found at index 3 in gen.

The Ultimate Guide to Slicing in Python

Method 4: Use a Generator and a For Loop

This example creates a Generator whose content is output to the terminal until a specific number is found.

def my_func():
    yield 10
    yield 20
    yield 30

gen = my_func()

for item in gen:
    if item == 20:
        print(item)
        break

This first line of the above code declares the function my_func(). This function will return, via a yield statement, a number (one/iteration).

Next, an object is declared and saved to gen.

If output to the terminal, an object similar to that shown below will display.

<generator object my_func at 0x0000022DE93D59A0>

The following line instantiates a for loop. This loop iterates through each yield statement in gen until the item contains a value of 20.

This value is output to the terminal, and the loop terminates via the break statement.

20

Summary

This article has provided four (4) ways to retrieve a single element from a Generator to select the best fit for your coding requirements.

Good Luck & Happy Coding!

Programming Humor – Python

“I wrote 20 short programs in Python yesterday. It was wonderful. Perl, I’m leaving you.” — xkcd