Закрытие соединения с `with` заявление

Question

Я хочу иметь класс, представляющий соединение по протоколу IMAP и использовать его с with заявление следующим образом:

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) 

Естественно это не удается, так как IMAPConnections не имеет атрибута close. Как я могу сохранить соединение и передать его функции __exit__, когда закончен оператор with?

Answer 1

Вам необходимо сохранить соединение в атрибутах объекта. Что-то вроде этого:

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() 

Вы также хотели бы реализовать list метод для вас класса.

Редактировать: Я только сейчас понял, какова ваша фактическая проблема. когда вы делаете with SomeClass(*args, **kwargs) as cc - это не значение, возвращаемое __enter__ методом. c - пример SomeClass. Это корень вашей проблемы, с которой вы вернули соединение от __enter__, и предположил, что указанное соединение - c.

Answer 2

Вам необходимо реализовать функцию __exit__() в пределах вашего класса IMAPConnection.

__enter__() функция вызывается перед выполнением кода в with блоке, где, как __exit__() называется при выходе из with блока.

Ниже представлена ​​структура выборки:

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

Проверил: Explaining Python's 'enter' and 'exit' для получения дополнительной информации.