Python with...as for custom context manager

Question

I wrote a simple context manager in Python for handling unit tests (and to try to learn context managers):

class TestContext(object):     test_count=1     def __init__(self):         self.test_number = TestContext.test_count         TestContext.test_count += 1      def __enter__(self):         pass      def __exit__(self, exc_type, exc_value, exc_traceback):         if exc_value == None:             print 'Test %d passed' %self.test_number         else:             print 'Test %d failed: %s' %(self.test_number, exc_value)         return True 

If I write a test as follows, everything works okay.

test = TestContext() with test:    print 'running test %d....' %test.test_number    raise Exception('this test failed') 

However, if I try to use with...as, I don't get a reference to the TestContext() object. Running this:

with TestContext() as t:     print t.test_number 

Raises the exception 'NoneType' object has no attribute 'test_number'.

Where am I going wrong?

Answer 1

Assuming that you need to access the context manager created in the with statement, __enter__ needs to return self. If you don't need to access it, __enter__ can return whatever you would like.

The with statement will bind this method’s return value to the target(s) specified in the as clause of the statement, if any.

This will work.

class TestContext(object):     test_count=1     def __init__(self):         self.test_number = TestContext.test_count         TestContext.test_count += 1      def __enter__(self):         return self      def __exit__(self, exc_type, exc_value, exc_traceback):         if exc_value == None:             print 'Test %d passed' % self.test_number         else:             print 'Test %d failed: %s' % (self.test_number, exc_value)         return True