Search code examples
pythonwith-statementcontextmanager

Surprising behavior of with keyword in python


I wanted to modify context manager behavior of an existing instance of a class (say, a database connection object). My initial idea was to monkey-patch __enter__ and __exit__ on the instance. To my surprise, that did not work. Monkey-patching the class achieves the desired effect (with a caveat that I am not sure that updating __class__ is a good idea).

What is the reason for this behavior of the with keyword? Essentially, I am looking for an explanation of why I should not be surprised. I could not find how the with is implemented, and I did not get the answer by reading PEP 343.

A runnable piece of code to illustrate.

import types


class My:
    def __enter__(self):
        print('enter')

    def __exit__(self, a, b, c):
        print('exit')


def add_behavior_to_context_manager(c):  # Does not work

    c_enter = c.__enter__
    c_exit = c.__exit__

    def __enter__(self):
        print('enter!')
        c_enter()
        return c

    def __exit__(self, exc_type, exc_value, exc_tb):
        c_exit(exc_type, exc_value, exc_tb)
        print('exit!')

    c.__enter__ = types.MethodType(__enter__, c)
    c.__exit__ = types.MethodType(__exit__, c)

    return c


def add_behavior_by_modifying_class(c):  # Works

    class MonkeyPatchedConnection(type(c)):
        def __enter__(self):
            print('enter!')
            return super().__enter__()

        def __exit__(wrapped, exc_type, exc_value, exc_tb):
            super().__exit__(exc_type, exc_value, exc_tb)
            print('exit!')

    c.__class__ = MonkeyPatchedConnection

    return c


my = add_behavior_to_context_manager(My())
print('Methods called on the instance of My work as expected: ')
my.__enter__()
my.__exit__(None, None, None)

print('Instance methods are ignored by the "with" statement: ')
with add_behavior_to_context_manager(My()):
    pass

print('Instead, class methods are called by the "with" statement: ')
with add_behavior_by_modifying_class(My()):
    pass

And the output:

Methods called on the instance of My work as expected: 
enter!
enter
exit
exit!
Instance methods are ignored by the "with" statement: 
enter
exit
Instead, class methods are called by the "with" statement: 
enter!
enter
exit
exit!

Solution

  • This is not specific to __enter__ and __exit__, but happens for other special methods as well. See https://docs.python.org/3/reference/datamodel.html#special-method-lookup:

    For custom classes, implicit invocations of special methods are only guaranteed to work correctly if defined on an object’s type, not in the object’s instance dictionary. [...]

    The rationale behind this behaviour lies with a number of special methods such as __hash__() and __repr__() that are implemented by all objects, including type objects. If the implicit lookup of these methods used the conventional lookup process, they would fail when invoked on the type object itself