Utilisation conviviale d'un Python itérable sur une séquence de gestionnaires de contexte

Question

J'écris une petite bibliothèque qui tente de fournir une file d'attente permanente pour l'envoi de tâches. Mon code de persistance fournit un moyen d'itérer sur les descriptions de travail en attente; Je voudrais également garantir que les tâches envoyées finiront par être marquées comme terminées ou échouées.

Pour ce faire, je d'abord mis en oeuvre pour que mon utilisateur peut faire:

for c in some_iterator_object: 
    with c as x: 
    ... 

Je n'aime pas cette solution pour plusieurs raisons. Tout d'abord, je veux récupérer une description de travail de ma file en une seule opération (et échouer si la file est vide), donc l'acquisition est faite par la méthode __next__ de l'itérateur, et la sortie dans le __exit__ du contexte directeur.

Pour que le gestionnaire de contexte soit appelé, mon __next__ renvoie une classe wrapper qui ne peut pas être substituée directement à la valeur, de sorte qu'il génère une erreur manifeste si l'utilisateur oublie d'appeler le gestionnaire de contexte.

Existe-t-il un moyen de réduire ces deux instructions en une seule? Idéalement, je voudrais laisser l'utilisateur faire tout

for x in some_iterator_object: 
    ... 

tout en étant capable d'intercepter des exceptions soulevées par le contenu du bloc pour.

EDIT: J'ai découvert en expérimentant que si je laisse un générateur inachevé obtenir déchets collectés, la déclaration de rendement déclenche une exception interne, donc je peux écrire quelque chose brut comme

try: 
    ... 
    success = False 
    yield val 
    success = True 
    ... 
finally: 
    if success: 
    ... 

Mais si je comprends bien , cela dépend du garbage collector pour fonctionner, et il semble être un mécanisme interne que je ne devrais pas vraiment toucher.

Answer 1

Si vous voulez que vos gestionnaires de contexte à saisir automatiquement car ils sont renvoyés par l'itérateur, vous pouvez écrire votre propre classe iterator comme ceci:

class ContextManagersIterator: 

    def __init__(self, it): 
     self._it = iter(it) 
     self._last = None 

    def __iter__(self): 
     return self 

    def __next__(self): 
     self.__exit__(None, None, None) 

     item = next(self._it) 
     item.__enter__() 
     self._last = item 

     return item 

    def __enter__(self): 
     return self 

    def __exit__(self, exc_type, exc_value, exc_traceback): 
     last = self._last 
     if last is not None: 
      self._last = None 
      return last.__exit__(exc_type, exc_value, exc_traceback) 

Exemple d'utilisation:

from contextlib import contextmanager 

@contextmanager 
def my_context_manager(name): 
    print('enter', name) 
    try: 
     yield 
    finally: 
     print('exit', name) 

sequence = [ 
    my_context_manager('x'), 
    my_context_manager('y'), 
    my_context_manager('z'), 
] 

with ContextManagersIterator(sequence) as it: 
    for item in it: 
     print(' work') 

# Output: 
# enter x 
# work 
# exit x 
# enter y 
# work 
# exit y 
# enter z 
# work 
# exit z 

Cette classe ContextManagersIterator prend soin d'appeler __enter__ sur ses valeurs juste avant qu'elles ne soient renvoyées. __exit__ est appelée juste avant qu'une autre valeur ne soit renvoyée (si tout s'est bien passé) ou lorsqu'une exception a été déclenchée dans la boucle.