try/except and context managers

I am trying to make my try/except blocks as small as possible. However, I don’t know how to reconcile this with a potentially lengthy with block:

try:
    with db.connect() as connection:
        # many lines interacting with the database
        # and more lines
        # and yet more lines

except ConnectionError:

Is there a way to write the try/except so it isn’t so long?

I suppose I could reorganize most of the code into a separate function:

try:
    with db.connect() as connection:
        do_stuff(connection)
except ConnectionError:

…but it seems like there should be a way to do it in fewer lines.

Answer

You don’t actually have to put the with statement in the try statement, if db.connect is the only thing that can raise the ConnectionError. If there is a connection error, there’s no connection to close.

try:
    connection = db.connect()
except ConnectionError:
    ...

with connection:
    ...

It’s not the call to db.connect itself that the with statement cares about; it’s the calls to connection.__enter__ and connection.__exit__.

Depending on what you do about the ConnectionError, the with statement should probably be guarded to avoid trying to use the undefined name connection.


If you don’t like the idea of being able to do something with connection between it being assigned and the with statement, you can use an ExitStack to enter the context immediately. For example,

from contextlib import ExitStack


with ExitStack() as es:
    try:
        connection = es.enter_context(db.connect())
    except ConnectionError:
        ...

    # do stuff with the connection