Error handling in context managers

[Peter Otten]

Israel Brewster wrote:

> I generally use context managers for my SQL database connections, so I can
> just write code like:
> 
> with psql_cursor() as cursor:
>     <do whatever>
> 
> And the context manager takes care of making a connection (or getting a
> connection from a pool, more likely), and cleaning up after the fact (such
> as putting the connection back in the pool), even if something goes wrong.
> Simple, elegant, and works well.
> 
> The problem is that, from time to time, I can't get a connection, the
> result being that cursor is None, and attempting to use it results in an
> AttributeError. So my instinctive reaction is to wrap the potentially
> offending code in a try block, such that if I get that AttributeError I
> can decide how I want to handle the "no connection" case. This, of course,
> results in code like:
> 
> try:
> with psql_cursor() as cursor:
> <do whatever>
> except AttributeError as e:
> <handle no-connection case>
> 
> I could also wrap the code within the context manager in an if block
> checking for if cursor is not None, but while perhaps a bit clearer as to
> the purpose, now I've got an extra check that will not be needed most of
> the time (albeit a quite inexpensive check).

That seems to be the cleanest approach.

> The difficulty I have with either of these solutions, however, is that
> they feel ugly to me - and wrapping the context manager in a try block
> almost seems to defeat the purpose of the context manager in the first
> place - If I'm going to be catching errors anyway, why not just do the
> cleanup there rather than hiding it in the context manager?
> 
> Now don't get me wrong - neither of these issues is terribly significant
> to me. I'll happily wrap all the context manager calls in a try block and
> move on with life if that it in fact the best option. It's just my gut
> says "there should be a better way", so I figured I'd ask: *is* there a
> better way? Perhaps some way I could handle the error internally to the
> context manager, such that it just dumps me back out? Of course, that
> might not work, given that I may need to do something different *after*
> the context manager, depending on if I was able to get a connection, but
> it's a thought. Options?

The problem is that there is no way to skip the with-block. It would be nice 
if you could raise a special exception in the __enter__() method that would 
achieve that.

For the time being you could 

(1) abuse a for loop

def psql_cursor():
    try:
      yield make_cursor()
    except NoConnection:
      pass

for cursor in psql_cursor():
    ...

or you could 

(2) return something that hopefully raises an exception soon, e. g.

$ cat conditional_context.py
import sys
from contextlib import contextmanager

class Cursor:
    def execute(self, sql):
        print("EXECUTING", sql)

class Exit(Exception):
    pass

class FailingCursor:
    def __getattr__(self, name):
        raise Exit

@contextmanager
def cursor():
    try:
        yield FailingCursor() if "--fail" in sys.argv else Cursor()
    except Exit:
        print("no connection")
        

with cursor() as cs:
    cs.execute("insert into...")
        
              
$ python3 conditional_context.py
EXECUTING insert into...
$ python3 conditional_context.py --fail
no connection

Option (1) works, but looks wrong while option (2) looks better, but does 
not work reliably in the general case, e. g. when you perform some costly 
calculation before the first attribute access.