Surprising behavior of with keyword in python

Question

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!

I'm not an expert on this, but in the PEP, the illustrative code in this section is interesting: https://peps.python.org/pep-0343/#specification-the-with-statement. Note the `exit = type(mgr).__exit__` (as opposed to `exit = mgr.__exit__`), and the similar behaviour for `__enter__`. — slothrop, Jun 16 '23 at 18:32
Yeah, I saw that - so at least we can be reasonably sure this behavior is not by accident. But there is no explanation of why the method is set on `type(mgr)`. — user443854, Jun 16 '23 at 19:39

score 1 · Accepted Answer · answered Jun 24 '23 at 00:05

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

Thank you. So, because `__enter__`, `__exit__` are "special", they are resolved in the special way. Something about it is not entirely satisfactory. For example, I struggle to come up with an example when the type object would be used in a context manager. — user443854, Jun 24 '23 at 01:04

Surprising behavior of with keyword in python

1 Answers1