Chaining context managers in Python

Let’s say we have a context manager that accepts two arguments:

@contextmanager
def colored(idx, color):
    push_style_color(idx, *color)
    yield
    pop_style_color()

How do we write a context manager that “freezes” the first argument?

with colored_text(pink):  # idx=COLOR_TEXT
    ...

One way is to bind the first argument:

def colored_text(color):
    return color(COLOR_TEXT, color)

But this function no longer looks like a context manager.

Because the wrapped function of a context manager is a generator, it’s possible to chain context managers in the way we chain generators:

@contextmanager
def colored_text(color):
    yield from color.__wrapped__(COLOR_TEXT, color)

This extra thought process gives us additional capability. Since we are writing a context manager again, we can extend it as a context manager:

@contextmanager
def color_replaced_if(cond, idx, color):
    if cond:
        yield from colored.__wrapped__(idx, color)
    else:
        yield