skyking
skyking

Reputation: 14395

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

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?

Upvotes: 5

Views: 1580

Answers (3)

hl037_
hl037_

Reputation: 3887

A softer version than the one proposed by Martineau, a little bit less polemic :

import sys

class CustomModule(sys.modules[__name__].__class__):
  """
  Custom module
  """
  def __enter__(self):
    print('enter')

  def __exit__(self, *args, **kwargs):
    print('exit')


sys.modules[__name__].__class__ = CustomModule

Instead of replacing the module (wich may cause countless problems), just replace the class by one inheriting from the original class. This way, the original module object is kept, there is no need for another ref (preventing garbage collection), and it will work with any custom importer. Note the important fact that a module object is created and added to the sys.modules BEFORE the code of the module is executed.

Note that using this way, you can add any magic method

Upvotes: 1

martineau
martineau

Reputation: 123491

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.

Upvotes: 2

Martijn Pieters
Martijn Pieters

Reputation: 1123400

__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).

Upvotes: 7

Related Questions