Reputation: 9432
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?
Upvotes: 43
Views: 12289
Reputation: 4725
Some of the code here is python 2 rather than 3, and didn't work when I ran. From https://realpython.com/python-with-statement/#writing-a-sample-class-based-context-manager the following example worked for me:
>>> class HelloContextManager:
... def __enter__(self):
... print("Entering the context...")
... return "Hello, World!"
... def __exit__(self, exc_type, exc_value, exc_tb):
... print("Leaving the context...")
... print(exc_type, exc_value, exc_tb, sep="\n")
...
>>> with HelloContextManager() as hello:
... print(hello)
...
Entering the context...
Hello, World!
Leaving the context...
None
None
None
Upvotes: 0
Reputation: 9986
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
Upvotes: 44
Reputation: 91099
def __enter__(self):
return self
will make it work. The value returned from this method will be assigned to the as
variable.
See also the Python doc:
If a target was included in the
with
statement, the return value from__enter__()
is assigned to it.
If you only need the number, you can even change the context manager's logic to
class TestContext(object):
test_count=1
def __init__(self):
self.test_number = TestContext.test_count
TestContext.test_count += 1
def __enter__(self):
return self.test_number
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
and then do
with TestContext() as test_number:
print test_number
Upvotes: 14
Reputation: 4755
According to PEP 343, the with EXPR as VAR
statement doesn't assign to VAR
the result of EXPR
, but rather the result of EXPR.__enter__()
. The first example worked because you referenced the test
variable itself.
Upvotes: 6