Brian McFarland
Brian McFarland

Reputation: 9432

Getting "NoneType object has no attribute" when using "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?

Upvotes: 43

Views: 12289

Answers (4)

baxx
baxx

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

Morgan Thrapp
Morgan Thrapp

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

glglgl
glglgl

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

art-solopov
art-solopov

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

Related Questions