Itérateurs, générateurs et itertools

Vous utilisez des for depuis le début du cursus. Mais comment fonctionnent-ils en interne ? Comment Python sait-il parcourir une liste, un dict, un fichier ? La réponse : le protocole d'itération — un mécanisme qui permet d'écrire ses propres objets itérables, de produire des données à la demande, et de traiter des flux infinis. C'est la clé de la performance pour de gros volumes.

Itérables vs itérateurs

Deux notions distinctes à ne pas confondre.

Itérable

Un itérable est un objet sur lequel on peut boucler avec for. Implémente __iter__() qui renvoie un itérateur.

Exemples d'itérables : list, tuple, str, dict, set, range, objets fichier...

Itérateur

Un itérateur est un objet qui produit les éléments un à un. Implémente __next__().

# Une liste est un itérable
liste = [1, 2, 3]

# iter() obtient un itérateur depuis un itérable
it = iter(liste)
print(type(it))

# next() avance d'un élément
print(next(it))
print(next(it))
print(next(it))

# Quand l'itérateur est épuisé, next() lève StopIteration
it = iter([1, 2])
next(it); next(it)
next(it)     # StopIteration

Ce que fait un for en coulisses

for x in liste:
    print(x)

...est équivalent à :

it = iter(liste)
while True:
    try:
        x = next(it)
    except StopIteration:
        break
    print(x)

Python cache cette complexité, mais c'est utile à connaître pour comprendre certains comportements.

Créer son itérable (protocole manuel)

On peut créer un objet itérable en implémentant __iter__ et __next__. Rarement nécessaire en pratique (les générateurs le font mieux), mais important à comprendre pour le TOSA.

class Compteur:
    """Itère de 1 à max (inclus)."""

    def __init__(self, max):
        self.max = max
        self.courant = 0

    def __iter__(self):
        return self        # l'objet est lui-même son itérateur

    def __next__(self):
        if self.courant >= self.max:
            raise StopIteration
        self.courant += 1
        return self.courant

c = Compteur(5)
for x in c:
    print(x, end=" ")

Les générateurs avec yield

Le yield transforme une fonction en générateur. Au lieu de calculer et retourner toute la liste, elle produit les éléments un par un, à la demande.

Premier exemple

def compter_jusqu_a(n):
    """Générateur des entiers de 1 à n."""
    i = 1
    while i <= n:
        yield i        # produit i et SUSPEND la fonction
        i += 1

# Usage comme n'importe quel itérable
for x in compter_jusqu_a(5):
    print(x, end=" ")

# On peut aussi le transformer en liste
print(list(compter_jusqu_a(5)))

Comprendre yield

yield n'est pas return. Quand Python rencontre yield :

1. Il produit la valeur au code appelant.
2. Il suspend la fonction à cet endroit.
3. Quand on appelle next(), la fonction reprend exactement là où elle s'était arrêtée.

def demo():
    print("  [début]")
    yield 1
    print("  [après 1er yield]")
    yield 2
    print("  [après 2e yield]")
    yield 3
    print("  [fin]")

g = demo()                  # ne lance RIEN encore
print("→ appel de next :")
print(next(g))              # exécute jusqu'au 1er yield
print("→ appel de next :")
print(next(g))
print("→ appel de next :")
print(next(g))

Voyez : la fonction reprend son exécution à chaque next, et les print intermédiaires s'exécutent au bon moment. C'est fondamentalement différent d'une fonction classique.

Gros avantage : économie mémoire

Une liste de 10 millions d'éléments occupe plusieurs centaines de Mo. Un générateur qui les produit à la demande : quelques centaines d'octets.

import sys

# Liste : tous les éléments en mémoire
liste_millions = [x ** 2 for x in range(10_000_000)]
print(f"Liste : {sys.getsizeof(liste_millions):,} octets")

# Générateur : quasi rien en mémoire
gen_millions = (x ** 2 for x in range(10_000_000))
print(f"Générateur : {sys.getsizeof(gen_millions):,} octets")

Expressions génératrices (rappel)

On a vu les expressions génératrices au chapitre 1 de la Partie 3 :

# Parenthèses = générateur, pas liste
carres = (x ** 2 for x in range(10))
print(carres)
print(list(carres))

Elles sont équivalentes à une fonction générateur courte :

def carres_gen():
    for x in range(10):
        yield x ** 2