6. Iterátory a generátory

Obsah

Protokol iterování (Iteration protocol)

Na předchozích semináři jsme mohli vidět použití protokolu iterování v cyklu for.

for x in [1, 2, 3, 4]:
    print(x)

Díky protokolu iterování funguje cyklus for (a všechny nástroje využívající postupující zleva doprava např. list comprehensions - ty uvidíme později) pro všechny objekty, které jsou iterovatelné (takzvané iterable).

Koncept iterable je generalizací pojmu sekvence. Objekt je považován za iterovatelný pokud je fyzicky uložen jako sekvence nebo se jedná o objekt, který během iterace produkuje jednu hodnotu v jeden okamžik (například během for loopu).

Terminologie: iterable vs iterator
iterator = iter([1, 2, 3])

next(iterator) # 1
next(iterator) # 2
next(iterator) # 3

next(iterator) # exception StopIteration

Pro podporu funkce iter() je nutné implementovat dunder metodu __iter__. Iterátor produkovaný funkcí iter() musí implementovat dunder metody __next__ a vyvolat StopIteration na konci produkce hodnot.

S iterable (iterovatelný objekt) jsme se již setkali mnohokrát:

point = {"x": 10, "y": 20}
point.keys() # view object - není to iterátor, stejně tak .items(), .values()

# iterátor lze ale jednoduše získat
iterator = iter(point)

next(iterator) # "x"
next(iterator) # "y"

next(iterator) # exception StopIteration

Rovněž zde vidíme důvod proč bylo nutné použít list() pro zobrazení celého výsledku range().

list(range(5)) # [0, 1, 2, 3, 4]

iterator = iter(range(5)) 

next(iterator) # 0
next(iterator) # 1
# ...

Stejně tak pro enumerate().

list(enumerate(range(5))) # [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4)]

iterator = iter(enumerate(range(5)))

next(iterator) # (0, 0)
next(iterator) # (1, 1)
# ...

Funkce zip(), enumerate(), filter() vrací iterable a rovněž iterable přijímají, lze je tedy zanořovat.

list(zip(enumerate(range(10)), range(10)))
# [((0, 0), 0),
#  ((1, 1), 1),
#  ((2, 2), 2),
#  ((3, 3), 3),
#  ((4, 4), 4),
#  ((5, 5), 5),
#  ((6, 6), 6),
#  ((7, 7), 7),
#  ((8, 8), 8),
#  ((9, 9), 9)]

Comprehension

Comprehension spolu s cykly jsou nejčastější případy použití iteračního protokolu.

Pokud budeme chtít umocnit seznam čísel, můžeme napsat klasický for cyklus.

squares = []

for number in numbers:
    squares.append(number ** 2)

Použití list comprehension:

squares = [number ** 2 for number in numbers]

Obecně platí předpis:

new_list = [expression for member in iterable]

Výsledkem list comprehension je vždy nový seznam. Comprehension je tedy dobré používat pouze pokud chceme vytvořit nový seznam hodnot. Dále pak platí, že pokud je list comprehension delší než dva řádky, je vhodnější (čitelnější) použít klasický for cyklus.

squares_minus = [number ** 2 for number in [number - 5 for number in numbers]]

Podmínky v list comprehension

Filtrování

Jeden ze způsobů použití podmínek v comprehension je filtrování. Nejprve se podíváme jak bychom napsali řešení klasicky:

sentence = 'the rocket came back from mars'

vowels = []

for char in sentence:
    if char in 'aeiou':
        vowels.append(char)

vowels   # ['e', 'o', 'e', 'a', 'e', 'a', 'o', 'a']

A nyní řešení pomoci list comprehension:

vowels = [i for i in sentence if i in 'aeiou']

Obecný předpis je tedy:

new_list = [expression for member in iterable (if conditional)]
Úprava prvků

Dále lze podmínky použít ke změně hodnot v seznamu. Nejprve klasické řešení:

numbers = [10, 20, -5, 10]
abs_numbers = []

for number in numbers:
    if number > 0:
        abs_numbers.append(number)
    else:
        abs_numbers.append(abs(number))

abs_numbers   # [10, 20, 5, 10]

A nyní řešení pomoci list comprehension:

numbers = [10, 20, -5, 10]

abs_numbers = [number if number > 0 else abs(number) for number in numbers]

Obecný předpis je tedy:

new_list = [expression (if conditional) for member in iterable]

Zanořování list comprehension

List comprehension můžeme zanořovat, situaci demonstrujeme výpočtem kartézského součinu. Nejprve klasickým způsobem:

colors = ['red', 'green', 'blue']
sizes = ['S', 'M', 'L', 'XL']

tshirts = []

for color in colors:
    for size in sizes:
        tshirts.append((color, size))

A nyní řešení pomoci list comprehension:

# všimněme si, že na pořadí samozřejmné záleží

tshirts_by_color = [(color, size) for color in colors for size in sizes]
# [('red', 'S'),
#  ('red', 'M'),
#  ...
#  ('blue', 'XL')]

tshirts_by_size = [(color, size) for size in sizes for color in colors]
# [('red', 'S'),
#  ('green', 'S'),
#  ...
#  ('blue', 'XL')]

assert tshirts_by_color != tshirts_by_size

Obecně pak můžeme použít předpis:

[expression for target1 in iterable1 if condition1
            for target2 in iterable2 if condition2 ...
            for targetN in iterableN if conditionN]

Set comprehension

Jak jsme viděli, v případě list comprehension je výsledkem vždy seznam. Pokud požadujeme aby výsledkem byla množina, můžeme použít set comprehension.

sentence = 'the rocket came back from mars'

unique_vowels = {i for i in sentence if i in 'aeiou'}
# {'a', 'e', 'o'}

Dict comprehension

Podobně pak v případě slovníku dict comprehension.

sentence = 'the rocket came back from mars'

word_len = {word: len(word) for word in sentence.split(' ')}
# {'the': 3, 'rocket': 6, 'came': 4, 'back': 4, 'from': 4, 'mars': 4}

Poznámka k rozsahu platnosti

V případě comprehension nejsou lokální proměnné dostupné:

[x for x in range(10)]

# nedostupná proměnná
x

# oproti tomu klasický for cyklus
output = []

for x in range(10):
    output.append(x)

# dostupná proměnná
x

Poznámka k výkonu

Ve většině případů comprehensions přinesou znatelné zrychlení (často až dvojnásobné). Iterace comprehension jsou v rámci interpretu prováděny rychlostí jazyka C (narozdíl od běžného for cyklu).

Především pro případy, kdy iterujeme přes velká data, je vždy lepší použít comprehension, pozor však na čitelnost kódu.

Generátorové funkce a výrazy

V novějších verzích jazyka Python je "prokrastinace" nebo také odložený výpočet využíván mnohem častěji než dříve. Jak se odložený výpočet projevuje? Namísto produkování celého výsledku najednou, jsou jednotlivé prvky (například prvky seznamu) produkovány v čase přístupu.

big_data = range(100000000000000)
small_data = range(1)

# 180 ns ± 14.4 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)
%timeit zip(small_data, small_data)
# 184 ns ± 13.5 ns per loop (mean ± std. dev. of 7 runs, 10000000 loops each)
%timeit zip(big_data, big_data)

big_zip = zip(big_data, big_data)
# jednotlivé hodnoty jsou generovány postupně
# zde se vypočítá první hodnota
next(big_zip)
# zde se vypočítá druhá hodnota
next(big_zip)

# nelze, tato hodnota neexistuje (generatory obecne nejsou indexovatelné)
big_zip[100000]

Generátorové funkce

Narozdíl od normálních funkcí, které skončí výpočet a vrátí hodnotu (pomoci příkazu return), generátorové funkce automaticky uspávají a probouzejí jejich vykonávání (pomoci příkazu yield).

Generátorové funkce jsou úzce spojené s iteračním protokolem. Aby bylo možné iterační protokol používat, jsou funkce obsahující yield kompilovány jako generátory. Nejedná se tedy o klasické funkce. Jejich architektura zaručuje, aby vracely objekt podporující iterační protokol. Později, při volání, vrací generátorový objekt, který podporuje iterační rozhraní (automaticky vytvoří metodu __next__).

V rámci generátorové funkce můžeme použít rovněž příkaz return, ten pak slouží na předčasné ukončení výpočtu (interně pomoci StopIteration).

# klasická funkce
def calculate_squares(n):
    result = []
    for i in range(n):
        result.append(i ** 2)
    return result


# generátorová funkce
def generate_squares(n):
    for i in range(n):
        yield i ** 2


# 749 μs ± 3.67 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit calculate_squares(20000)

# 944 μs ± 2.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)
%timeit list(generate_squares(20000))

Generátorové funkce je nejvhodnější používat tehdy, když nevíme, zda budeme potřebovat celý výsledek nebo pouze jeho část.

# v rámci for loopu
for square in generate_squares(20000000):
    print(square)
    if square > 100:
        break

# jako generátor
generator = generate_squares(20000000)

next(generator) # 0
next(generator) # 1
next(generator) # 4
# ...

Generátory jsou single-iteration objekty. To znamená, že je přes ně možné iterovat pouze jedenkrát. Pokud se chceme k výsledkům vracet, je nutné je ukládat.

generator = generate_squares(20)

for square in generator:
    print(square)

# zde se již nic nevypíše, generátor je prázdný
for square in generator:
    print(square)

Pokud výsledky uložíme do seznamu, můžeme je procházet vícekrát.

generator = list(generate_squares(20))

for square in generator:
    print(square)

# zde se vypise
for square in generator:
    print(square)

Generátorové výrazy

Obdoba list comprehension v kontextu iterátorů. Generátory můžeme vytvářet jednodušeji než definováním funkce.

# klasicky list comprehension (výpočet proběhne okamžitě)
[number ** 2 for number in range(10)]
# vrací - [0, 1, 4, 9, 16, 25, 36, 49, 64, 81]

# generátorový výraz
(number ** 2 for number in range(10))
# vrací - <generator object <genexpr> at 0x106b73ac0>

generator = (number ** 2 for number in range(10))

# zde se vypočítá první hodnota
next(generator)
# zde se vypočítá druhá hodnota
next(generator)

Příkaz yield from

Ve verzi 3.3 byl přidán příkaz yield fromna delegování generátoru. Pokud máme nějaký generátor a chceme z něj vytvořit další generátor je možné napsat funkci:

def my_gen(gen):
    yield from gen

Příkaz yield from je rovněž elegantním způsobem na vytvoření generátoru (iterátoru) z uložené sekvence.

def g(x):
    yield from range(x, 0, -1)
    yield from range(x)


list(g(5))
# vysledek: [5, 4, 3, 2, 1, 0, 1, 2, 3, 4]

Sekvence vs Iterátory

Nyní se podíváme na vytvoření vlastní sekvence pomoci implementování dunder metod __len__ a __getitem__.

class Sentence:
    """Represents one sentence."""

    def __init__(self, text):
        self.text = text
        self.words = text.split(' ')
    
    def __repr__(self):
        return f"Sentence({self.text})"
    
    def __len__(self):
        return len(self.words)
    
    def __getitem__(self, index):
        return self.words[index]


sentence = Sentence("Ahoj svete jak se mas")

# magie! Sentence je iterable díky dunder metodě __getitem__
for word in sentence:
    print(word)

# Ahoj
# svete
# jak
# se
# mas    

sentence[2] # "jak"

len(sentence) # 5

Nyní se podíváme na stejnou věc pohledem iterátoru. Zde je důležité rozlišit iterable a iterator. Třída Sentence je iterable, třída SentenceIterator je pak iterátor, který vrací metoda Sentence.__iter__.

Pro vytvoření iterátoru je tedy nutné implementovat dunder metody __iter__ a __next__.

class Sentence:
    """Represents one sentence."""

    def __init__(self, text):
        self.text = text
        self.words = text.split(' ')
    
    def __repr__(self):
        return f"Sentence({self.text})"
    
    def __iter__(self):
        return SentenceIterator(self.words)


class SentenceIterator():
    def __init__(self, words):
        self.words = words
        self.index = 0
    
    def __next__(self):
        try:
            word = self.words[self.index]
        except IndexError:
            raise StopIteration
        
        self.index += 1
        return word
    
    def __iter__(self):
        return self
    

sentence = Sentence("Ahoj svete jak se mas")

# magie! nyní díky dunder metodě __iter__
for word in sentence:
    print(word)

len(sentence)

Další možností je situaci implementovat jako generátor. Tady pouze upozorním, že nevyužíváme opravdového postupného výpočtu, pro demonstraci je však tento příklad dostačující.

class Sentence:
    """Represents one sentence."""

    def __init__(self, text):
        self.text = text
        self.words = text.split(' ')
    
    def __repr__(self):
        return f"Sentence({self.text})"
    
    def __iter__(self):
        yield from self.words


sentence = Sentence("Ahoj svete jak se mas")

# magie
for word in sentence:
    print(word)

iterator = iter(sentence)

next(iterator) # "Ahoj"
next(iterator) # "svete"
next(iterator) # "jak"
next(iterator) # "se"
next(iterator) # "mas"

next(iterator) # StopIteration

Jak využít toho, aby byl výpočet vyhodnocován postupně? Zkusme navrhnout generátor opravdový.

def split_generator(text, sep=[" "]):
    """Creates generator for splitted text by given separator"""
    word = []

    for char in text:
        if char in sep:
            if word:
                yield "".join(word)
                word = []
        else:
            word.append(char)

    if word:
        yield "".join(word)


class Sentence:
    """Represents one sentence."""
    
    def __init__(self, text):
        self.text = text
    
    def __repr__(self):
        return f"Sentence({self.text})"
    
    def __iter__(self):
        yield from split_generator(self.text):
        
        # zde je rovněž možné použít generator expression
        # return (word for word in split_generator(self.text))

        # nejlepším řešením je však vrátit samotný generátor
        # return split_generator(self.text)


sentence = Sentence("Ahoj svete jak se mas")

# magie
for word in sentence:
    print(word)
    if word == "svete":
        break 

# Ahoj
# svete


iterator = iter(sentence)

next(iterator) # "Ahoj"
next(iterator) # "svete"
next(iterator) # "jak"
next(iterator) # "se"
next(iterator) # "mas"

next(iterator) # StopIteration

Funkce vracející iterátory

Funkce any a all

Vrací True v případě že je některý prvek/jsou všechny prvky iterable vyhodnoceny na pravdu. Pokud je iterable prázdná, vrací False.

any([[], [1, 2, 3], [1]]) # True

all([[], [1, 2, 3], [1]]) # False

Funkce map

Vrací iterátor (pozor nikoli sekvenci), který aplikuje funkci na každý prvek předaného iterable a tyto výsledky jsou vráceny pomoci yield.

map(lambda x: x ** 2, [1, 2, 3, 4])

Funkce map může přijímat libovolný počet iterable, v takovém případě však musí aplikovaná funkce přijímat stejný počet argumentů jako je počet iterable.

import operator


# operator.mul přijímá 2 argumenty
map(operator.mul, [1, 2, 3, 4], [2, 3, 4, 5])

Funkce filter

Vrací iterátor (pozor nikoli sekvenci) vytvořený z prvků iterable pro které se funkce vyhodnotí na pravdu. Pokud je funkce rovna None jsou z iterable tímto způsobem odstraněny všechny prvky vyhodnocené na nepravdu.

filter(None, [[], [1, 2, 3], [1]])

filter(lambda x: x % 3, range(20))

Funkce filter(function, iterable) je ekvivalentní generátorovému výrazu:

# pokud je function rovno None
(item for item in iterable if function(item))

# jinak
(item for item in iterable if item)

List comprehension vs map/filter

Funkci map/filter je vhodné nahrazovat list comprehension, můžeme dosáhnout vyššího výkonu. Jedná se však o konstantní overhead volání funkce.

import timeit


# 10
print(timeit.timeit("list(map(lambda x: x * x, range(10)))", number=10000))
# 0.016546586997719714

print(timeit.timeit("[x * x for x in range(10)]", number=10000))
# 0.008891337001841748

# 10000
print(timeit.timeit("list(map(lambda x: x * x, range(10000)))", number=10000))
# 7.573531197998818

print(timeit.timeit("[x * x for x in range(10000)]", number=10000))
# 5.103068967000581

# 100000
print(timeit.timeit("list(map(lambda x: x * x, range(100000)))", number=10000))
# 96.5402136729972

print(timeit.timeit("[x * x for x in range(100000)]", number=10000))
# 63.06499578600051

Modul itertools

Modul implementující sadu užitečných iterátorů. Tyto iterátory preferujeme nad vlastní implementací z důvodu rychlosti a paměťové úspornosti.

itertools.starmap

Vytvoří iterátor, který vyhodnocuje předanou funkci s argumenty získanými z iterable. Narozdíl od klasického map přijímá jeden iterable, ten však může obsahovat n-tice (v případě funkce přijímající n argumentů). Sémanticky podobné použítí * při volání funkce.

import itertools
import operator


# operator.mul přijímá 2 argumenty
itertools.starmap(operator.mul, [[1, 2], [2, 3], [3, 4], [4, 5]])

itertools.filterfalse

Pracuje podobně jako filter, výsledný iterátor však produkuje prvky pro které je funkce nepravda.

import itertools


itertools.filterfalse(None, [[], [1, 2, 3], [1]])

itertools.filterfalse(lambda x: x % 3, range(20))

Ostatní iterátory

Celkový výčet iterátorů dostupných v balíčku itertools je dostupný zde.