Python multiple context managers in one class

Question

I would like to be able to write code like this:

with obj.in_batch_mode:
    obj.some_attr = "some_value"
    obj.some_int = 142
    ...

when I want obj to wait with sending updates about itself until multiple jobs are completed. I have hooks on __setattr__ that take some time to run, and the changes can be sent together.

I do not want to use code like this, since it increases the risk of forgetting to leave batch_mode (which is what the with keyword is good for):

obj.enter_batch_mode()
obj.some_attr = "some_value"
obj.some_int = 142
...
obj.exit_batch_mode()

I have not been able to figure out how to implement this. Just typing with obj: (and simply implementing with on obj) does not read anywhere near as descriptive.

You need to implement the enter and exit functions in order to use the built in with statement — omri_saadon
– omri_saadon, Commented Jul 19, 2015 at 21:14
I want to implement multiple context managers on different properties so i can write with obj.in_batch_mode and with obj.in_verbose_mode etc. instead of the with obj example above, which does not read well — Drathier
– Drathier, Commented Jul 19, 2015 at 21:21
Could you have some mode state in the object that affects the way the object behaves in the __enter__ and __exit__ methods? It wouldn't be quite as explicit as it sounds like you want since the mode would not necessarily be included in your with obj: call... — Engineero
– Engineero, Commented Jul 19, 2015 at 21:24

Dunes · Accepted Answer · 2015-07-22 18:38:55Z

Generally, a very simple way to implement context managers is to use the contextlib module. Writing a context manager becomes as simple as writing a single yield generator. Before the yield replaces the __enter__ method, the object yielded is the return value of __enter__, and the section after the yield is the __exit__ method. Any function on your class can be a context manager, it just needs the be decorated as such. For instance, take this simple ConsoleWriter class:

from contextlib import contextmanager

from sys import stdout
from io import StringIO
from functools import partial

class ConsoleWriter:

    def __init__(self, out=stdout, fmt=None):
        self._out = out
        self._fmt = fmt

    @property
    @contextmanager
    def batch(self):
        original_out = self._out
        self._out = StringIO()
        try:
            yield self
        except Exception as e:
            # There was a problem. Ignore batch commands.
            # (do not swallow the exception though)
            raise
        else:
            # no problem
            original_out.write(self._out.getvalue())
        finally:
            self._out = original_out

    @contextmanager
    def verbose(self, fmt="VERBOSE: {!r}"):
        original_fmt = self._fmt
        self._fmt = fmt
        try:
            yield self
        finally:
            # don't care about errors, just restore end
            self._fmt = original_fmt

    def __getattr__(self, attr):
        """creates function that writes capitalised attribute three times"""
        return partial(self.write, attr.upper()*3)


    def write(self, arg):
        if self._fmt:
            arg = self._fmt.format(arg)
        print(arg, file=self._out)

Example usage:

writer = ConsoleWriter()
with writer.batch:
    print("begin batch")
    writer.a()
    writer.b()
    with writer.verbose():
        writer.c()
    print("before reentrant block")
    with writer.batch:
        writer.d()
    print("after reentrant block")
    print("end batch -- all data is now flushed")

Outputing:

begin batch
before reentrant block
after reentrant block
end batch -- all data is now flushed
AAA
BBB
VERBOSE: 'CCC'
DDD

Pynchia · Accepted Answer · 2015-07-19 22:03:56Z

1

If you are after a simple solution and do not need any nested mode-change (e.g. from STD to BATCH to VERBOSE back to BATCH back to STD)

class A(object):
    STD_MODE = 'std' 
    BATCH_MODE = 'batch'
    VERBOSE_MODE = 'verb'

    def __init__(self):
        self.mode = self.STD_MODE

    def in_mode(self, mode):
        self.mode = mode
        return self

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self.mode = self.STD_MODE

obj = A()
print obj.mode
with obj.in_mode(obj.BATCH_MODE) as x:
    print x.mode
print obj.mode

outputs

std
batch
std

answered Jul 19, 2015 at 22:03

Pynchia

11.7k5 gold badges38 silver badges49 bronze badges

1 Comment

Drathier Over a year ago

this could easily be modified to allowing multiple modes by using a list of modes or a bitmap

Drathier · Accepted Answer · 2015-07-21 19:52:06Z

This builds on Pynchia's answer, but adds support for multiple modes and allows nesting of with statements, even with the same mode multiple times. It scales O(#nested_modes) which is basically O(1).

Just remember to use stacks for data storage related to the modes.

class A():
    _batch_mode = "batch_mode"
    _mode_stack = []

    @property
    def in_batch_mode(self):
        self._mode_stack.append(self._batch_mode)
        return self

    def __enter__(self):
        return self

    def __exit__(self, type, value, tb):
        self._mode_stack.pop()
        if self._batch_mode not in self._mode_stack:
            self.apply_edits()

and then I have these checks wherever I need them:

if self._batch_mode not in self._mode_stack:
    self.apply_edits()

It is also possible to use methods for modes:

with x.in_some_mode(my_arg):

just remember to save my_arg in a stack within x, and to clear it from the stack when that mode is popped from the mode stack.

The code using this object can now be

with obj.in_batch_mode:
    obj.some_property = "some_value"

and there are no problems with nesting, so we can add another with obj.in_some_mode: wherever without any hard-to-debug errors or having to check every function called to make sure the object's with-statements are never nested:

def b(obj):
    with obj.in_batch_mode:
        obj.some_property = "some_value"

x = A()
with x.in_batch_mode:
    x.my_property = "my_value"
    b(x)

eran · Accepted Answer · 2015-07-19 21:34:25Z

Maybe something like this:

Implement helper class

class WithHelperObj(object):
    def __init__(self,obj):
        self.obj = obj

    def __enter__(self):
        self.obj.impl_enter_batch()

    def __exit__(self, exc_type, exc_value, traceback):
        self.obj.impl_exit_batch()

    class MyObject(object):
        def in_batch_mode(self):
            return WithHelperObj(self)

In the class itself, implement method instead of field, to use with the with statement

    def impl_enter_batch(self):
        print 'In impl_enter_batch'

    def impl_exit_batch(self):
        print 'In impl_exit_batch'

    def doing(self):
        print 'doing'

Then use it:

o = MyObject()
with o.in_batch_mode():
    o.doing()

Collectives™ on Stack Overflow

Python multiple context managers in one class

4 Answers 4

Comments

1 Comment

Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

4 Answers 4

Comments

1 Comment

Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related