Context manager for SMTP connections

Question

My code right now look like this

@contextmanager
def smtp_connection(host, user=None, passwd=None, timeout=5):
    conn = None # smell here
    try:
        conn = SMTP(host=host, timeout=timeout)
        conn.ehlo_or_helo_if_needed()
        if user and passwd:
            conn.login(user=user, password=passwd)
        logger.debug('SMTP connected')
        yield conn
    except Exception as e:
        raise e
    finally:
        if conn: # and here
            conn.quit()

In recipes for ExitStack there is a suggestion to replace try-finally and flag variables with this

with ExitStack() as stack:
    stack.callback(cleanup_resources)
    result = perform_operation()
    if result:
        stack.pop_all()

But this doesn't use result in cleanup_resources. So in my case it still would be

with ExitStack() as stack:
    result = None
    stack.callback(lambda conn: conn.quit())
    result = POP3() # code from above here
    if result:
        stack.pop_all()

Answer 1

except Exception as e:
    raise e

First, you could simplify this to

except Exception:
    raise

Second, since you only re-raise it, you are not prepared to handle any Exception that could appear either with your code or the code managing the connection returned by your context manager. Thus you don't need that except clause.

Now the only thing left to manage is the state of the connection. Since there is no exception handling performed by your code, you are free to create the connection out of the try … finally and use that mechanism to only close the connection whatever happened:

@contextmanager
def smtp_connection(host, user=None, passwd=None, timeout=5):
    conn = SMTP(host=host, timeout=timeout)
    try:
        conn.ehlo_or_helo_if_needed()
        if user and passwd:
            conn.login(user=user, password=passwd)
        logger.debug('SMTP connected')
        yield conn
    finally:
        conn.quit()

---

You could also simplify the design by separating concerns. As it stand, your function does two things:

• manage the connection;
• perform some initial setup.

By delegating to a second context manager, you could separate these two behaviors into reusable bits:

@contextmanager
def smtp_connection(host, timeout):
    connection = SMTP(host=host, timeout=timeout)
    try:
        yield connection
    finally:
        connection.quit()

@contextmanager
def smtp_setup(host, user=None, passwd=None, timeout=5):
    with smtp_connection(host, timeout) as conn:
        conn.ehlo_or_helo_if_needed()
        if user and passwd:
            conn.login(user=user, password=passwd)
        logger.debug('SMTP connected')
        yield conn

But, looking at the second context manager, there is nothing to manage anymore since there is no teardown/cleanup anymore. Thus it is best to provide it as an independant function:

@contextmanager
def smtp_connection(host, timeout=5):
    connection = SMTP(host=host, timeout=timeout)
    try:
        yield connection
    finally:
        connection.quit()

def smtp_setup(conn, user=None, passwd=None):
    conn.ehlo_or_helo_if_needed()
    if user and passwd:
        conn.login(user=user, password=passwd)
    logger.debug('SMTP connected')

You would then need to change your calling code from

with smtp_connection(h, u, p, t) as conn:
    # do stuff

to

with smtp_connection(h, t) as conn:
    smtp_setup(conn, u, p)
    # do stuff

---

In the end, I personally don't like to have to manage the try ... finally inside the context manager. I don't find that natural as I preffer to write an explicit class using __enter__ and __exit__:

class smtp_connection:
    def __init__(self, host, timeout):
        self.smtp = SMTP(host=host, timeout=timeout)
    def __enter__(self):
        return self.smtp
    def __exit__(self, exc_type, exc_value, exc_trace):
        self.smtp.quit()

Using a class could also let you wrap this thing around SMTP directly:

class smtp_connection(SMTP):
    def __init__(self, host, timeout):
        super().__init__(host=host, timeout=timeout)
    def __enter__(self):
        return self
    def __exit__(self, exc_type, exc_value, exc_trace):
        self.quit()
    def setup(self, user=None, password=None):
        self.ehlo_or_helo_if_needed()
        if user and passwd:
            self.login(user=user, password=passwd)
        logger.debug('SMTP connected')

Use it like:

with smtp_connection(h, t) as conn:
    conn.setup(u, p)
    # do stuff