Can we mix contextmanager decorator with __enter__() and __exit__() methods in another class inside-CodePudding

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.

CodePudding user response：

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)