Closing a connection with a `with` statement

Question

I want to have a class that represents an IMAP connection and use it with a with statement as follows:

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)

Naturally this fails since IMAPConnections has no attribute close. How can I store the connection and pass it to the __exit__ function when the with statement is finished?

Darth Kotik · Accepted Answer · 2016-09-23T07:36:19.133

You need to store connection in object attributes. Something like this:

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

You would also want to implement list method for you class.

Edit: I just now realized what your actual problem is. when you do with SomeClass(*args, **kwargs) as c c is not a value returned by __enter__ method. c is instance of SomeClass. That's a root of your problem you returned connection from __enter__ and assumed that c is said connection.

I added some more explanation in my answer. Hope it's helpfull — Darth Kotik, Sep 23 '16 at 07:36

score 0 · Answer 2 · edited May 23 '17 at 10:27

You need to implement a __exit__() function within your IMAPConnection class.

__enter__() function is called before executing the code within the with block where as __exit__() is called while exiting the with block.

Below is the sample structure:

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

Check: Explaining Python's 'enter' and 'exit' for more information.

Closing a connection with a `with` statement

2 Answers2