Why can't a module be a context manager (to a 'with' statement)?

Question

Suppose we have the following mod.py:

def __enter__():
    print("__enter__<")

def __exit__(*exc):
    print("__exit__< {0}".format(exc))

class cls:
    def __enter__(self):
        print("cls.__enter__<")

    def __exit__(self, *exc):
        print("cls.__exit__< {0}".format(exc))

and the following use of it:

import mod

with mod:
    pass

I get an error:

Traceback (most recent call last):
  File "./test.py", line 3, in <module>
    with mod:
AttributeError: __exit__

According to the documentation the documentation the with statement should execute as follows (I believe it fails at step 2 and therefore truncate the list):

1. The context expression (the expression given in the with_item) is evaluated to obtain a context manager.
2. The context manager’s __exit__() is loaded for later use.
3. The context manager’s __enter__() method is invoked.
4. etc...

As I've understood it there is no reason why __exit__ could not be found. Is there something I've missed that makes a module not able to work as a context manager?

Answer 1

__exit__ is a special method, so Python looks it up on the type. The module type has no such method, which is why this fails.

See the Special method lookup section of the Python datamodel documentation:

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.

Note that this applies to all special methods. For example, if you added a __str__ or __repr__ function to a module it'll not be called when printing the module, either.

Python does this to make sure type objects are hashable and representable too; if Python didn't do this then trying to put a class object into a dictionary would fail when a __hash__ method was defined for that class (as that method would expect an instance to be passed in for self).

Answer 2

You can't do it easily for the reasons stated in @Martijn Pieters answer. However with a little extra work it is possible, because the values in sys.modules don't have to be instances of the built-in module class, they can be instances of your own custom class with the special methods a context manager requires.

Here's applying that to what you want to do. Given the following mod.py:

import sys

class MyModule(object):
    def __enter__(self):
        print("__enter__<")

    def __exit__(self, *exc):
        print("__exit__> {0}".format(exc))

# replace entry in sys.modules for this module with an instance of MyModule
_ref = sys.modules[__name__]
sys.modules[__name__] = MyModule()

And the following use of it:

import mod

with mod:
    print('running within context')

Will produces this output:

__enter__<
running within context
__exit__> (None, None, None)

See this question for information about why the _ref is needed.