[REQUEST] Improved uninstall for rich.traceback.install and rich.pretty.install

[antonymilne]

We are very excited to have recently adopted rich on the Kedro framework. In short, our setup is:

logging.config.dictConfig(logging_config)  # where logging_config uses rich logging handler
rich.traceback.install(show_locals=True, suppress=[click])
rich.pretty.install()

Some people have reported various issues (e.g. #2455, #2408), and more generally, users have asked whether they can uninstall rich tracebacks. See also #1947.

Now there are a few problems:

• the above code is baked into the framework and not accessible to Kedro users. Hence a user cannot easily capture the original exception handler that is returned by rich.traceback.install in order to restore it
• when running on IPython, rich.traceback.install modifies various things like ip._showtraceback. However, the function only returns sys.excepthook. Hence, even if a user were able to capture the call as old_excephook = rich.traceback.install(), it would not provide sufficient information to fully undo the effects of rich.traceback.install()
• rich.pretty.install has caused fewer problems for users so is less important to try to uninstall, but it does not return anything at all and hence, even if you had access to the code that called rich.pretty.install, there's no way to reliably reverse the call

Suggested solutions

A new rich.traceback.uninstall functionality that fully reverses the effect of rich.traceback.install. This would not require the user to be able to access the call to install in the first place and would also work on IPython. Similarly for rich.pretty.uninstall (but less important to us).

Current workarounds

# Undo rich.traceback.install
import sys
sys.excepthook = sys.__excepthook__

# In IPython
from IPython.core.interactiveshell import InteractiveShell
ip._showtraceback = InteractiveShell()._showtraceback
ip.showtraceback = InteractiveShell().showtraceback
ip.showsyntaxerror = InteractiveShell().showsyntaxerror

# Undo rich.pretty.install
sys.displayhook = sys.__displayhook__ 

# In IPython
from IPython.core.formatters import PlainTextFormatter
ip = get_ipython()
ip.display_formatter.formatters["text/plain"] = PlainTextFormatter()

This is not quite right because it doesn't restore things to how they were before the rich installs; it just resets them to the Python/IPython defaults. Hence on platforms such as Databricks, where it's difficult to figure out what the correct in-built Databricks customisations to exception handling etc., the above uninstalls aren't correct because they don't restore the settings to their databricks pre-rich-install state.

[willmcgugan]

Rather than uninstall it, could you provide an env var to not install it in the first place?

[antonymilne]

Thanks for the very rapid response! That is indeed our currently proposed solution (probably wouldn't a an environment variable due to how kedro works, but some sort of user-specified option that would skip the call to the rich installs). It's not ideal from our perspective for various reasons (e.g. it's not clear where this would live and how it would operate differently for both IPython and non-IPython workflows). Also, given that rich.traceback.install already sort of has a mechanism built-in to undo its replacement of sys.excepthook, I thought that a fully functioning uninstall feature would belong on rich also. So I just thought I'd gauge appetite for doing this here and whether it's something you'd accept a PR for or not.

[willmcgugan]

I'd accept a PR for an uninstall function.

[antonymilne]

Cool, thank you! Any tips on implementation? e.g. how would you store the pre-install IPython settings. A global variable? Make a class to contain the state and then install/uninstall functions that access that class?

[noklam]

It needs to handle both cases:

1. IPython - It is not using sys.excepthook and I haven't figured out how to restore it.
2. Non IPython - this seems to be quite straightforward as the install() return the original hook, you can just override it with sys.exceptionhook = old_hook