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

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