Brian McFarland Brian McFarland - 19 days ago 7
Python Question

Python with...as for custom context manager

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

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
Comments