Fermeture d'une connexion avec un `with` déclaration

Question

Je veux avoir une classe qui représente une connexion IMAP et l'utiliser avec une déclaration with comme suit:

class IMAPConnection: 
    def __enter__(self): 
     connection = imaplib.IMAP4_SSL(IMAP_HOST) 

     try: 
      connection.login(MAIL_USERNAME, MAIL_PASS) 
     except imaplib.IMAP4.error: 
      log.error('Failed to log in') 

     return connection 

    def __exit__(self, type, value, traceback): 
     self.close() 

with IMAPConnection() as c: 
    rv, data = c.list() 
    print(rv, data) 

Naturellement cela échoue puisque IMAPConnections n'a pas d'attribut close. Comment puis-je stocker la connexion et la passer à la fonction __exit__ lorsque l'instruction with est terminée?

Answer 1

Vous devez stocker la connexion dans les attributs d'objet. Quelque chose comme ceci:

class IMAPConnection: 
    def __enter__(self): 
     self.connection = imaplib.IMAP4_SSL(IMAP_HOST) 

     try: 
      self.connection.login(MAIL_USERNAME, MAIL_PASS) 
     except imaplib.IMAP4.error: 
      log.error('Failed to log in') 

     return self.connection 

    def __exit__(self, type, value, traceback): 
     self.connection.close() 

Vous voudrait également implémenter la méthode list pour vous classe. Edit: Je viens de réaliser maintenant quel est votre problème actuel. lorsque vous faites with SomeClass(*args, **kwargs) as cc n'est pas une valeur renvoyée par la méthode __enter__. c est l'instance de SomeClass. C'est une racine de votre problème que vous avez renvoyé la connexion de __enter__ et supposé que c est ladite connexion.

Answer 2

Vous devez implémenter une fonction __exit__() dans votre classe IMAPConnection. Fonction

__enter__() est appelé avant d'exécuter le code dans le bloc with où que __exit__() est appelé en sortant du bloc with.

est inférieure à la structure de l'échantillon:

def __exit__(self, exc_type, exc_val, exc_tb): 
    # Close the connection and other logic applicable 
    self.connection.close() 

Vérifier: Explaining Python's 'enter' and 'exit' pour plus d'informations.