Is there a __enter__ and __exit__ equivalent for all functions in a class?

Question

One functionality of python that I found very handy when working with databases (or files) are the __enter__ and __exit__ functions you can give to a class. Now by using the with statement you can make sure that in this block __enter__ is first called (and you can open the database or file) and after it's done __exit__ is called (and you can close a database or file.

I want to open and close a sqlite transaction every time a function from my Database class is called. I can do it at the start and end of every function, but since it has to be done for every function is that class, I was wondering, are there methods that get called before and after each function call? Like SetUp and TearDown in unittesting.

Answer 1

You can decorate every member function with a pie decorator, something like

@transaction
def insertData(self):
    # code

and transaction is a decorator you define to wrap the function with a pre and post. Yes, you have to do it for every function. Here is an example

def transaction(f):
    def pre():
        print "pre transaction"
    def post():
        print "post transaction"

    def wrapped(*args):
        pre()
        f(*args)
        post()

    return wrapped


class Foo(object):
    def __init__(self):
        print "instantiating"

    def doFoo(self):
        print "doing foo"

    @transaction
    def doBar(self, value):
        print "doing bar "+str(value)

@transaction
def foofunc():
    print "hello"

foofunc()

f=Foo()
f.doFoo()
f.doBar(5)

.

stefanos-imac:python borini$ python decorator.py 
pre transaction
hello
post transaction
instantiating
doing foo
pre transaction
doing bar 5
post transaction

The alternative is that you use a metaclass, like this:

import types


class DecoratedMetaClass(type):
    def __new__(meta, classname, bases, classDict):
        def pre():
            print "pre transaction"
        def post():
            print "post transaction"
        newClassDict={}
        for attributeName, attribute in classDict.items():
            if type(attribute) == types.FunctionType:
                def wrapFunc(f):
                    def wrapper(*args):
                        pre()
                        f(*args)
                        post()
                    return wrapper
                newAttribute = wrapFunc(attribute)
            else:
                newAttribute = attribute
            newClassDict[attributeName] = newAttribute
        return type.__new__(meta, classname, bases, newClassDict)



class MyClass(object):

    __metaclass__ = DecoratedMetaClass

    def __init__(self):
        print "init"
    def doBar(self, value):
        print "doing bar "+str(value)
    def doFoo(self):
        print "doing foo"



c = MyClass()
c.doFoo()
c.doBar(4)

This is pure black magic, but it works

stefanos-imac:python borini$ python metaclass.py
pre transaction
init
post transaction
pre transaction
doing foo
post transaction
pre transaction
doing bar 4
post transaction

You normally don't want to decorate the __init__, and you may want to decorate only those methods with a special name, so you may want to replace

        for attributeName, attribute in classDict.items():
            if type(attribute) == types.FunctionType:

with something like

        for attributeName, attribute in classDict.items():
            if type(attribute) == types.FunctionType and "trans_" in attributeName[0:6]:

This way, only methods called trans_whatever will be transactioned.