Feature/unload ipython extension

[nikolas-claussen]

Summary

This PR adds the %unload_ext jaxtyping line magic to unload the jaxtyping extension in a Jupyter notebook. This way, users can disable type checking without having to restart their notebook kernels. This is useful for executing cells where getting type annotations is not practical.

Changes

• jaxtyping/_ipython_extension.py:

  • Adds function unload_ipython_extension

• jaxtyping/__init__.py:

  • Imports unload_ipython_extension

• test/test_ipython_extension.py

  • Adds function test_unload_extension_disables_typechecking

[patrick-kidger]

Oh yikes, this is a bit terrifying. If this is what it takes to perform an unregister then I'm afraid this isn't something I feel comfortable supporting!

[nikolas-claussen]

Reasonable! I tracked down where exactly the loaded extension is stored to make the unloading more targeted. See my latest commit, or below. Let me know what you think.

def unload_ipython_extension(ipython):
    """
    Support `%unload_ext jaxtyping` to remove the jaxtyping AST transformer
    and unregister the `%jaxtyping.typechecker` magic.
    """
    # Names registered via @line_magic use the explicit string we provided.
    extension_name = "jaxtyping.typechecker"
    try:
        # 0) Disable runtime typechecking globally (covers already-decorated functions).
        config.jaxtyping_disable = True
        
        # 1) Remove any JaxtypingTransformer from the AST transformers.
        ipython.ast_transformers = [
            t for t in getattr(ipython, "ast_transformers", None)
            if not isinstance(t, JaxtypingTransformer)]

        # 2) Unregister the `%jaxtyping.typechecker` line magic. 
        mm = getattr(ipython, "magics_manager", None)
        magics = getattr(mm, "magics", None)
        magics["line"].pop(extension_name)
    except Exception:
        RuntimeError("Failed to unload jaxtyping.typechecker magic")