How to chain context managers in python?

Question

Long story short, what would be the right way to get second snippet of code to work exactly like the first one?

stack_device = [None]
stack_context = [None]

@contextlib.contextmanager
def device(device):
    stack_device.append(device)
    try:
        yield
    finally:
        stack_device.pop()


@contextlib.contextmanager
def context(ctx):
    stack_context.append(ctx)
    try:
        with device("dev"):
            yield
    finally:
        stack_context.pop()


with context("myctx"):
    print(stack_device[-1])  # -> dev
    print(stack_context[-1]) # -> ctx

And that one, of course, would not have the right device set when I need it:

stack_device = [None]
stack_context = [None]

class Device():
    def __init__(self, device):
        self.device = device

    def __enter__(self):
        stack_device.append(self.device)
        return

    def __exit__(self, type, value, traceback):
        stack_device.pop()


class Context():
    def __init__(self, ctx):
        self.ctx = ctx

    def __enter__(self):
        with Device("cls_dvc"):
            stack_context.append(self.ctx)
            return

    def __exit__(self, type, value, traceback):
        stack_context.pop()


with Context("myctx"):
    print(stack_device[-1])  # -> None !!!
    print(stack_context[-1]) # -> myctx

What would be the right way to achieve same behaviour in the second case as in first case?

Answer 1

The key to the failure of your code is that a with statement calls enter when its block, and exit when the block ends. Having a with Device block inside the enter of Context means that the return statement, even though inside the block, leaves it, thus triggering the exit of Device. You can see this course of action by adding prints inside of each special method.

There are several possible solutions to make it work:

1. If the context managers are independent, like in Robert Kearns' solution, instead of creating a super-context to create both your contexts you can use the standard Python chaining to create them in order, from left to right:

   with Context("myctx"), Device("cls_dvc"):

2. If you need access to Device attributes inside Context class (the only reason to chain them the way you do), then again you have two solutions:
   • Redesign your class structure so that a single __enter__ initiates both contexts and a single __exit__ cleans both contexts. In your design, only Context needs to be implemented as context manager, Device can be a regular class.
   • If you need to keep Device as an independent context manager (e.g., to use it by itself somewhere else), then Jack Taylor's solution is what you need. Note that there are several cases in Python (for example, open("filename")) that can be used as both "regular" and as "context managers". In this case, all the logic is in regular methods, while __enter__ and __exit__ only contain calls to the regular methods, so you don't have to call the special methods directly, as in the provided answers

Answer 2

You need to create a Device object inside your Context class, call the Device object's __enter__ method in the Context __enter__ method, and call the Device object's __exit__ method in the Context __exit__ method. If there is an error, then you can either handle it in the Context __exit__ method or the Device __exit__ method, whichever is more appropriate.

stack_device = [None]
stack_context = [None]

class Device:
    def __init__(self, device):
        self.device = device

    def __enter__(self):
        stack_device.append(self.device)
        return self

    def __exit__(self, err_type, err_value, traceback):
        stack_device.pop()


class Context:
    def __init__(self, ctx):
        self.ctx = ctx
        self.cls_dvc = Device("cls_dvc")

    def __enter__(self):
        self.cls_dvc.__enter__()
        stack_context.append(self.ctx)
        return self

    def __exit__(self, err_type, err_value, traceback):
        stack_context.pop()
        self.cls_dvc.__exit__(err_type, err_value, traceback)


with Context("myctx"):
    print(stack_device[-1])  # -> cls_dvc
    print(stack_context[-1]) # -> myctx

Answer 3

I get the right output by putting the with Device() manager inside the with Context().

stack_device = [None]
stack_context = [None]

class Device():
    def __init__(self, device):
        self.device = device

    def __enter__(self):
        stack_device.append(self.device)
        return

    def __exit__(self, type, value, traceback):
        stack_device.pop()


class SubContext():
    def __init__(self, ctx):
        self.ctx = ctx

    def __enter__(self):
        stack_context.append(self.ctx)
        return


    def __exit__(self, type, value, traceback):
        stack_context.pop()

class Context:

    def __init__(self, ctx):
        self.ctx = SubContext(ctx)
        self.device = Device('dev')

    def __enter__(self):
        self.ctx.__enter__()
        self.device.__enter__()

    def __exit__(self, type, value, traceback):
        self.ctx.__exit__(type, value, traceback)
        self.device.__exit__(type, value, traceback)

with Context("myctx"):
    print(stack_device[-1])
    print(stack_context[-1])