In python3.8 I’m very familiar with the traditional __enter__
and __exit__
magic methods but new to the @contextlib.contextmanager
decorator. Is it possible to mix the two patterns inside a single with
statement?
The following (highly contrived) script should explain the problem more clearly. Is there a definition of ContextClass.enter_context_function()
and ContextClass.exit_context_function()
(I imagine something needs to change inside __init__
as well) that only use the context_function()
function and makes the unit tests pass? Or are these patterns mutually exclusive?
import contextlib NUMBERS = [] @contextlib.contextmanager def context_function(): NUMBERS.append(3) try: yield finally: NUMBERS.append(5) class ContextClass: def __init__(self): self.numbers = NUMBERS self.numbers.append(1) def __enter__(self): self.numbers.append(2) self.enter_context_function() # should append 3 return self def __exit__(self, exc_type, exc_val, exc_tb): self.exit_context_function() # should append 5 self.numbers.append(6) def function_call(self): self.numbers.append(4) def enter_context_function(self): # FIX ME! pass def exit_context_function(self): # FIX ME! pass if __name__ == "__main__": import unittest class TestContextManagerFunctionAndClass(unittest.TestCase): def test_context_function_and_class(self): with ContextClass() as cc: cc.function_call() self.assertEqual(NUMBERS, [1, 2, 3, 4, 5, 6]) unittest.main()
I understand there are better ways to solve a similar problem (specifically rewriting context_function
as a class with its own __enter__
and __exit__
methods, but I’m trying to better understand exactly how the contextmanager decorator works.
Advertisement
Answer
No change in the __init__
is necessary. The manual way which “makes the unit tests pass” would be:
def enter_context_function(self): self._context_mgr = context_function() self._context_mgr.__enter__() def exit_context_function(self): self._context_mgr.__exit__(None, None, None)
However, it’s kind of missing the point of context-managers. They’re intended to be used in a with-statement.
Also note that, as written, the NUMBERS.append(5)
line (the “teardown”) may not be reached if the code after yielding raises. It should be written like this:
@contextlib.contextmanager def context_function(): NUMBERS.append(3) try: yield finally: NUMBERS.append(5)