Python с... как для настраиваемого менеджера контекста

Я написал простой диспетчер контекстов в Python для обработки модульных тестов (и для изучения контекстных менеджеров):

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

Если я напишу тест следующим образом, все будет хорошо.

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

Однако, если я пытаюсь использовать с... как, я не получаю ссылку на объект TestContext(). Выполнение этого:

with TestContext() as t:
    print t.test_number

Вызывает исключение 'NoneType' object has no attribute 'test_number'.

Где я иду не так?

Ответ 1

Предполагая, что вам нужно получить доступ к менеджеру контекста, созданному в инструкции with, __enter__ необходимо вернуть self. Если вам не нужен доступ, __enter__ может возвращать все, что вы хотели бы.

Оператор WITH свяжет эти методы с возвращаемым значением для целевых объектов, указанных в предложении as, если таковые имеются.

Это будет работать.

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

Ответ 2

def __enter__(self):
    return self

заставит его работать. Значение, возвращаемое этим методом, будет присвоено переменной as.

См. также документ Python:

Если цель была включена в оператор with, ему присваивается возвращаемое значение из __enter__().

Если вам нужен только номер, вы даже можете изменить логику менеджера контекста на

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

а затем do

with TestContext() as test_number:
    print test_number

Ответ 3

Согласно PEP 343, оператор with EXPR as VAR не присваивает VAR результат EXPR, а результат EXPR.__enter__(). Первый пример работал, потому что вы ссылались на переменную test.