How to override a python instance's __exit__ method

Question

I want to change the __exit__ method from within the __enter__ method in a python class. How can I do this? Here is what I tried:

class TestClass:
    def __enter__(self):
        func_type = type(self.__exit__)
        self.__exit__ = func_type(lambda *_: print("exit override"), self)

        # Also tried the next line
        # self.__exit__ = lambda *_: print("exit override")

        return self

    def __exit__(self, *args):
        print("exiting")

with TestClass() as test:
    pass

Later, I want to make that change conditional that's why I don't want to change the __exit__ method itself. Is there a way to do this other than test for the condition in the __exit__ method as well?

Answer 1

You can't do that with __exit__, because it's always taken from the class:

mgr = (EXPR)
exit = type(mgr).__exit__  <---- here
value = type(mgr).__enter__(mgr)

https://peps.python.org/pep-0343/#specification-the-with-statement

As suggested in the comments, the natural way to do that would be to attach a condition to a manager instance rather than changing the method:

class TestClass:
    def __enter__(self):
        if condition:
            self.whatever = xxx
        return self

    def __exit__(self, *args):
        if self.whatever == xxx:
            ...