CODEWITHSUNDEEP
pythonpython-3.xwith-statementclass-method
Richard Neumann
Richard Neumann

Reputation: 3371

__enter__ and __exit__ on class level in python 3

I am unsuccessfully trying to get the magic with-statement methods __enter__ and __exit__ running on class-level:

class Spam():

    @classmethod
    def __enter__(cls):
        return cls

    @classmethod
    def __exit__(cls, typ, value, tb):
        cls.cleanup_stuff()


with Spam:
    pass

However, this will result in an AttributeError:

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

Is it possible to use the __enter__ and __exit__ methods on class-level anyway?

Upvotes: 8

Views: 9831

Answers (3)

Micael Jarniac
Micael Jarniac

Reputation: 421

You can combine @classmethod and @contextlib.contextmanager to make a class method that acts as a context manager:

from __future__ import annotations

from contextlib import contextmanager
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from typing import Iterator, Self

class Spam:
    @classmethod
    @contextmanager
    def context(cls) -> Iterator[type[Self]]:
        try:
            # Setup here
            yield cls
        finally:
            cls.cleanup_stuff()

with Spam.context():
    pass

Note that the order of the decorators matters.

Reference:

Upvotes: 0

mata
mata

Reputation: 69082

__enter__ and __exit__ are special methods, and as such only work correctly when defined on a object's type, not in it's instance dictionary.

Now Spam is a instance of type, and type(Spam).__enter__ and type(Spam).__exit__ do not exist. Therefore you get an attribute error.

To make this work, the methods would need to be declared on the metaclass of the class you want to use. Example:

class Spam(type):

    def __enter__(cls):
        print('enter')
        return cls

    def __exit__(cls, typ, value, tb):
        print('exit')

class Eggs(metaclass=Spam):
    pass

with Eggs:
    pass

Now Eggs is an instance of Spam (type(Eggs) == Spam, and therefore type(Eggs).__enter__ and type(Eggs).__exit__ do exist).

However defining a metaclass just to use an instance of it as a context manager seems a little over the top. The more straight forward solution starting from your example would be to just use

with Spam():
    pass

Or if you want to reuse the same instance later:

spam = Spam()
with spam:
    pass

Upvotes: 16

myaut
myaut

Reputation: 11514

It seems that CPython doesn't call a bound method like instance.__exit__, it seeks over instance type, doing something like type(instance).__dict__['__exit__'] than calls it. And since type(Spam) is a special type object (not a Spam itself), it doesn't contain __exit__ method.

I tried to workaround that using metaclasses, but wasn't successful. __getattr__ doesn't work either.

See here: https://github.com/python/cpython/blob/2545fdbd4b4a6a77b132fccf816578f59b609be5/Objects/typeobject.c#L1362

  • Py_TYPE is similiar to type(self)
  • _PyType_LookupId walks over type(self).__dict__ (no __getattr__ call here)

Python 2 implementation is different, but main idea about getting type(self) applies to it too

Upvotes: 0

Related Questions