pythonpython-click

How to use Python Click's `ctx.with_resource` to capture tracebacks in (sub-)commands/groups


In Click, Context.with_resource's documentation states it can be used to:

[r]egister a resource as if it were used in a with statement. The resource will be cleaned up when the context is popped.

From this, I understand that the context manager that I pass to Context.with_resource will be exited after execution of the root CLI group and any of its sub-groups and commands. This seems to work fine with an example such as this, where I am redirecting stdout to a file:

import contextlib
import click

@contextlib.contextmanager
def my_redirect_stdout(file):
    with open(file, mode="w") as fp:
        with contextlib.redirect_stdout(fp):
            yield

@click.group()
@click.pass_context
@click.argument("file")
def cli(ctx, file):
    ctx.with_resource(my_redirect_stdout(file))

@cli.command()
def hello():
    print(f"this goes to a file")

if __name__ == "__main__":
    cli()

However, this does not work when I try to capture exception tracebacks in the following way:

import contextlib
import sys
import traceback
import click

@contextlib.contextmanager
def capture_traceback(file):
    with open(file, mode="w") as fp:       
        try:
            yield 
        except:
            print(f"exception!")
            traceback.print_exc(file=fp)
            sys.exit(1)

@click.group()
@click.pass_context
@click.argument("file")
def cli(ctx, file):
    ctx.with_resource(capture_traceback(file))

@cli.command()
def hello():
    raise ValueError("error!")

if __name__ == "__main__":
    cli()

The exception does not seem to be caught within the capture_traceback function. Instead, it is printed by the interpreter as usual. It seems as though Click is catching the error, closing the context manager, and then re-raising. How can I catch exceptions from any group or command, print the traceback to a file, and exit the program (without printing the traceback to the terminal/stderr)?


Solution

  • Click does not propagate up the errors like this. The click.BaseCommand class owns a click.Context which is storing these context managers in a simple ExitStack. When __exit__ is called on the context, it closes the exit stack without forwarding the error info. The code is roughly:

    class Context:
        ...
        def __exit__(self, exc_type, exc_value, tb):
            self._depth -= 1
            if self._depth == 0:
                self.close()
            pop_context()
        ...
    
        def close(self) -> None:
            self._exit_stack.close()
            self._exit_stack = ExitStack()
    

    Because click doesn't forward the exception info, you cant work on it. With that said, it does look like it's possible to override the context class used by the application.

    Disclaimer This following code is very brittle. It would probably be better to open an issue in the upstream repo requesting it. At the very least, you would want to pin the version of click to an exact version such that you can ensure this patch continues to function.

    ...
    
    class MyContext(click.Context):
        def __exit__(self, exc_type, exc_value, tb):
            self._depth -= 1
            if self._depth == 0:
                self._exit_stack.__exit__(exc_type, exc_value, tb)
                self._exit_stack = contextlib.ExitStack()
            click.core.pop_context()
    
    
    cli.context_class = MyContext
    
    
    if __name__ == "__main__":
        cli()