How do __enter__ and __exit__ work in Python decorator classes?

Question

I'm trying to create a decorator class that counts how many times a function is called, but I'm getting an error message that says:

    "TypeError: __exit__() takes exactly 1 argument (4 given)"

and I really don't know how I'm giving it four arguments. My code looks like this:

class fcount2(object):
    __instances = {}
    def __init__(self, f):
        self.__f = f
        self.__numcalls = 0
        fcount2.__instances[f] = self

    def __call__(self, *args, **kwargs):
        self.__numcalls += 1
        return self.__f(*args, **kwargs)

    def __enter__(self):
        return self

    def __exit__(self):
        return self

    @staticmethod
    def count(f):
        return fcount2.__instances[self.__f].__numcalls


@fcount2
def f(n):
    return n+2

for n in range(5):
    print f(n)   
print 'f count =',f.count

def foo(n):
    return n*n

with fcount2(foo) as g:
    print g(1)
    print g(2)
print 'g count =',g.count
print 'f count =',f.count

with fcount2(f) as g:
    print g(1)
    print g(2)
print 'g count =',g.count
print 'f count =',f.count

with f:
    print f(1)
    print g(2)
print 'g count =',g.count
print 'f count =',f.count

Are there some other parameters I should (or shouldn't) be passing into the def exit function? Any tips or ideas would be appreciated.

As an aside, my line of code that says "print 'f count =',f.count" appears to be outputting the memory address rather than the value, but that's a whole different problem.

Answer 1

the __exit__() method should accept information about exceptions that come up in the with: block. See here.

The following modification of your code works:

def __exit__(self, exc_type, exc_value, tb):
    if exc_type is not None:
        traceback.print_exception(exc_type, exc_value, tb)
        # return False # uncomment to pass exception through

    return True

Then you can try raising an exception in one of your with: blocks and it'll be caught in __exit__().