Easiest way to add a function to existing class

Question

I'm using the python's built-in shelve module to manage some simple dictionaries. The problem I'm having is I want to use with shelve.open(filename) as f:, but when I try it claims DbfilenameShelf has no attribute __exit__.

So, I'm guessing the easiest way to do this is to wrap it in another class and add an __exit__ function to that wrapper. I tried this:

class Wrapper(shelve.DbfilenameShelf):
    def __exit__(self):
        self.close()
    def __init__(self, filename, writeback=False):
        shelve.DbfilenameShelf.__init__(self, filename, flag='c', protocol=None, writeback=False)

But when I tried to instantiate the wrapper like so: wrapped = Wrapper(filename) it tells me I'm giving it an invalid argument.

Error as requested:

Traceback (most recent call last):
File "<input>", line 1, in <module>
File "<input>", line 5, in __init__
File "C:\Python27\Lib\shelve.py", line 223, in __init__
Shelf.__init__(self, anydbm.open(filename, flag), protocol, writeback)
File "C:\Python27\Lib\anydbm.py", line 85, in open
return mod.open(file, flag, mode)
File "C:\Python27\Lib\dbhash.py", line 18, in open
return bsddb.hashopen(file, flag, mode)
File "C:\Python27\Lib\bsddb\__init__.py", line 364, in hashopen
d.open(file, db.DB_HASH, flags, mode)
DBInvalidArgError: (22, 'Invalid argument')    

Answer 1

Don't subclass it. Python comes with a tool for automatically calling close(), contextlib.closing:

from contextlib import closing
with closing(shelve.open(filename)) as f:
    # your 'with' block here

will automatically call the close() method of the object returned by shelve.open(filename) at the end of the with block.

Answer 2

You're subclassing the wrong thing and missing an __enter__ method. You probably want this:

class contextShelf(shelve.shelve):
  def __enter__(self):
    return self

  def __exit__(self, exc_type, exc_value, exc_trace):
    self.close()

Because you're adding methods, but not changing the __init__ signature or adding any extra steps there's no reason you need to redefine __init__. The base class' __init__ will be called automatically.