Python - with statement with optional object

Question

I have a script that processes some data and if a database/file is present, writes some data into it. I specify the database or file as configargparse(argparse) argument. I need to clean (close file, db) in some organized way in case exceptions occur.

Here is my init:

import sqlite3
import confargparse
import sys
parser.ArgParser(...) 
parser.add('--database', dest='database',
    help='path to database with grabbers', metavar='FILE',
    type=lambda x: arghelper.is_valid_file(parser, x))
parser.add('-f', '--file', type=configargparse.FileType(mode='r'))
args = parser.parse_args()

I did it using if and try:

if args.database:
    conn = sqlite3.connect(args.database)
    c = conn.cursor()
# same init for file

try:
    while True: # do something, it might be moved to some main() function
        result = foo()
        if args.database:
            c.execute('Write to database {}'.format(result))
        # same
            # for file        
finally:
    if args.database:
        conn.close() 
    # same
        # for file
except KeyboardInterrupt:
    print 'keyboard interrupt'

Can it be done with the with statement? Something like (here comes ()?():() from C):

with ((args.database)?
      (conn = sqlite3.connect(args.database)):
      (None)) as db, same for file:

and then refer to the db inside the with clause and check if they exist?

Could you just let file default to os.devnull and db default to :memory:? — Holloway
– Holloway, Commented Feb 18, 2016 at 12:10
I should close file and db on the (unexpected) end. can I run `.close()' in such cases? — d21d3q
– d21d3q, Commented Feb 18, 2016 at 12:16
If you use the file path `/dev/null/', you can open, write to and close it as you would any other file. For the database, sqlite lets you create an 'in-memory' database that is lost when the execution ends. There may be memory issues if you have a lot of data; I'm not sure. — Holloway
– Holloway, Commented Feb 18, 2016 at 12:23

gil · Accepted Answer · 2016-02-18 14:03:51Z

To answer your question first. It can be done, using contextlib. But I'm not sure how much you would gain from this.

from contextlib import contextmanager

@contextmanager
def uncertain_conn(args):
    yield sqlite3.connect(args.database) if args.database else None

# Then you use it like this
with uncertain_conn(args) as conn:
    # conn will be the value yielded by uncertain_conn(args)
    if conn is not None:
        try:
            # ...

But as I said, while turning a generator function into a context manager is cool and personally I really like the contextmanager decorator, and it does give you the functionality you say you want, I don't know if it really helps you that much here. If I were you I'd probably just be happy with if:

if args.database:
    conn = sqlite3.connect(args.database)
    try:
       # ...

There are a couple of things you can simplify with with, though. Check out closing, also from contextlib (really simple, I'll just quote the documentation):

contextlib.closing(thing)

Return a context manager that closes thing upon completion of the block. This is basically equivalent to:
from contextlib import contextmanager

@contextmanager def closing(thing):
    try:
        yield thing
    finally:
        thing.close()

So the above code can become:

if args.database:
    conn = sqlite3.connect(args.database)
    with closing(conn):
        # do something; conn.close() will be called no matter what

But this won't print a nice message for KeyboardInterrupt. If you really need that, then I guess you still to have to write out the try-except-finally yourself. Doing anything more fanciful is probably not worth it. (And note that except must precede finally, otherwise you get a syntax error.)

And you can even do this with suppress (but requires a bit of caution; see below)

from contextlib import suppress

with suppress(TypeError):
    conn = sqlite3.connect(args.database or None)
    with closing(conn):
        # do business

with suppress(error): do_thing is equivalent to

try:
    do_thing
except error:
    pass

So if args.database evaluates to False, the second line is effectively connect(None), which raises a TypeError, which will be caught by the context manager and the code below will be skipped. But the risk is that it will suppress all TypeErrors in its scope, and you may not want that.

Alfe · Accepted Answer · 2016-02-18 12:46:28Z

You can create your own context manager in such cases. Create one which handles both connections. A context manager is a class which has methods __enter__() and __exit__(). One is called before entering the with clause, one is called when it is left (how ever).

Here's an example for how to do this in your case:

def f(cond1, cond2):

  class MultiConnectionContextManager(object):
    def __init__(self, cond1, cond2):
      self.cond1 = cond1
      self.cond2 = cond2
    def __enter__(self):
      print "entering ..."
      if self.cond1:
        # self.connection1 = open(...)
        print "opening connection1"
      if self.cond2:
        # self.connection1 = open(...)
        print "opening connection2"
      return self
    def __exit__(self, exc_type, exc_value, traceback):
      print "exiting ..."
      if self.cond1:
        # self.connection1.close()
        print "closing connection1"
      if self.cond2:
        # self.connection2.close()
        print "closing connection2"

  with MultiConnectionContextManager(cond1, cond2) as handle:
    if cond1:
      # handle.connection1.read()
      print "using handle.connection1"
    if cond2:
      # handle.connection2.read()
      print "using handle.connection2"

for cond1 in (False, True):
  for cond2 in (False, True):
    print "=====", cond1, cond2
    f(cond1, cond2)

You can call this directly to see the outcome. Replace the prints with your real statements for opening, using, and closing the connections.

Collectives™ on Stack Overflow

Python - with statement with optional object

2 Answers 2

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related