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.

Solution 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__().

Solution 2

The reason why you are getting that error is becuase __exit__() takes 3 arguments along with self. These are:

1. exception_type
2. exception_value
3. traceback

You can define __exit__() method two ways:

def __exit__(self, exception_type, exception_value, traceback):

or

def __exit__(self, *args, **kwargs):

with statement supports the concept of a runtime context defined by a context manager implemented using __enter__() & __exit__() method pair. Please visit this link for detailed info.

__enter__() method will enter the runtime context and return either this object or another object related to the runtime context.

__exit__() method will exit the runtime context and return a Boolean flag indicating if any exception that occurred should be suppressed. The values of these arguments contain information regarding the thrown exception. If the values equal to None means no exception was thrown.