CODEWITHSUNDEEP
pythonderived-classwith-statement
frans
frans

Reputation: 9798

Pythons 'with'-statement: correctly nest/derive classes with __enter__/__exit__

How should I correctly nest the with-related behavior of classes (e.g. when deriving or instantiating)?

This works for me but I wonder if there's a dedicated way to do it:

class class_a:
    def __init__(self):
        print('class_a::__init__')

    def __enter__(self):
        print('class_a::__enter__')
        return self

    def __exit__(self, type, exit, tb):
        print('class_a::__exit__')


class class_b(class_a):
    def __init__(self):
        class_a.__init__(self)
        print('class_b::__init__')

    def __enter__(self):
        class_a.__enter__(self)
        print('class_b::__enter__')
        return self

    def __exit__(self, type, exit, tb):
        class_a.__exit__(self, type, exit, tb)
        print('class_b::__exit__', type, exit, tb)

with class_b():
    print('ready')
    try:
        signal.pause()
    except:
        pass

One way to do this differently would be to implement class_b like this:

class class_b:
    def __init__(self):
        self._class_a_inst = class_a()
        print('class_b::__init__')

    def __enter__(self):
        self._class_a_inst.__enter__()
        print('class_b::__enter__')
        return self

    def __exit__(self, type, exit, tb):
        self._class_a_inst.__exit__(type, exit, tb)
        print('class_b::__exit__', type, exit, tb)

Is there any difference regarding the __enter__() / __exit__() behavior?

Upvotes: 3

Views: 471

Answers (1)

Steve Jessop
Steve Jessop

Reputation: 279355

Ideally, use contextlib.contextmanager. For the case of deriving:

import contextlib

class context_mixin:
    def __enter__(self):
         self.__context = self.context()
         return self.__context.__enter__()
    def __exit__(self, *args):
         return self.__context.__exit__(*args)

class class_a(context_mixin):
    @contextlib.contextmanager
    def context(self):
         print('class_a enter')
         try:
             yield self
         finally:
             print('class_a exit')

class class_b(class_a):
    @contextlib.contextmanager
    def context(self):
        with super().context():
            print('class_b enter')
            try:
                yield self
            finally:
                print('class_b exit')

In Python 2, super() needs to be super(class_b, self).

There is a change in behaviour compared with your code: this code exits b before exiting a, meaning that the scopes nest. You've written your code to do them in the other order, although that's easy enough to change. Often it makes no difference, but when it does matter you usually want things to nest. So for an (admittedly-contrived) example, if class_a represents an open file, and class_b represents some file format, then the exit path for class_a will close the file, while the exit path for class_b will write any buffered changes that have yet to be committed. Clearly b should happen first!

For the case of holding another object:

class class_b(context_mixin):
    def __init__(self):
        self.a = class_a()
    @contextlib.contextmanager
    def context(self):
        with self.a:
            print('class_b enter')
            try:
                yield self
            finally:
                print('class_b exit')

Upvotes: 3

Related Questions