Repository: patrick-kidger/jaxtyping (language: python)

Issue:
Applying `install_import_hook` to Script
I'm a big fan of jaxtyping and use it almost everywhere. However, I've found that it would be nice to apply `install_import_hook` directly to a self-contained script. For example:

```
from jaxtyping import install_script_hook

# This is like install_import_hook, but applies decorators to the functions in this self-contained script.
install_script_hook(("beartype", "beartype"))

# This now has runtime type checking.
def fn(x: int):
    pass

if __name__ == "__main__":
    fn(123)
```

Is this possible? The [documentation](https://docs.kidger.site/jaxtyping/api/runtime-type-checking/) suggests that there's no way other than splitting the script into an `entry_point.py` that calls `install_import_hook` and a `main.py` that contains the functions that import hook should be applied to.


Discussion / hints:
I think this should be possible with a bit of hackery:
```python
import inspect
import types
from collections.abc import Callable


def install_script_hook(*decorators: Callable):
    global_objs = inspect.stack()[1].frame.f_globals
    for name, value in global_objs.copy().items():
        if isinstance(value, types.FunctionType):
            for decorator in decorators:
                value = decorator(value)
            global_objs[name] = value

# Example usage:
#
# import beartype
#
# ... the entire rest of your script ...
#
# install_script_hook(beartype.beartype)  # called once all your functions are defined
#
# if __name__ == "__main__":
#    ...  # run stuff
```
Right now this will annotate all global functions. You could probably have it annotate other things too, if you're willing to start introspecting `__closure__` dictionaries and the like.
Thanks, this is exactly what I was looking for!

Fix patch:
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 68f9260..2c3fd73 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -30,6 +30,7 @@
 )
 from ._decorator import jaxtyped as jaxtyped
 from ._import_hook import install_import_hook as install_import_hook
+from ._ipython_extension import load_ipython_extension as load_ipython_extension
 
 
 # Now import Array and ArrayLike
diff --git a/jaxtyping/_ipython_extension.py b/jaxtyping/_ipython_extension.py
new file mode 100644
index 0000000..50fec53
--- /dev/null
+++ b/jaxtyping/_ipython_extension.py
@@ -0,0 +1,35 @@
+from ._import_hook import _JaxtypingTransformer
+
+
+try:
+    from IPython.core.magic import line_magic, Magics, magics_class
+
+    @magics_class
+    class ChooseTypecheckerMagics(Magics):
+        @line_magic("jaxtyping.typechecker")
+        def typechecker(self, typechecker):
+            # remove old _JaxtypingTransformer, if present
+            self.shell.ast_transformers = list(
+                filter(
+                    lambda x: not isinstance(x, _JaxtypingTransformer),
+                    self.shell.ast_transformers,
+                )
+            )
+
+            # add new one
+            self.shell.ast_transformers.append(
+                _JaxtypingTransformer(typechecker=typechecker)
+            )
+
+except ImportError:
+    pass
+
+
+def load_ipython_extension(ipython):
+    try:
+        ipython.register_magics(ChooseTypecheckerMagics)
+    except NameError:
+        raise NameError(
+            "ChooseTypecheckerMagics is not defined.\n\n"
+            + "You may be trying to use IPython extension without IPython installed."
+        )


Test patch:
diff --git a/test/requirements.txt b/test/requirements.txt
index 99e4780..0df2227 100644
--- a/test/requirements.txt
+++ b/test/requirements.txt
@@ -4,3 +4,4 @@ equinox
 jaxlib
 pytest
 typeguard<3
+IPython
diff --git a/test/test_ipython_extension.py b/test/test_ipython_extension.py
new file mode 100644
index 0000000..a42330d
--- /dev/null
+++ b/test/test_ipython_extension.py
@@ -0,0 +1,151 @@
+import pytest
+from IPython.testing.globalipapp import start_ipython
+
+from .helpers import ParamError
+
+
+@pytest.fixture(scope="session")
+def session_ip():
+    yield start_ipython()
+
+
+@pytest.fixture(scope="function")
+def ip(session_ip):
+    session_ip.run_cell(raw_cell="import jaxtyping")
+    session_ip.run_line_magic(magic_name="load_ext", line="jaxtyping")
+    session_ip.run_line_magic(
+        magic_name="jaxtyping.typechecker", line="beartype.beartype"
+    )
+    yield session_ip
+
+
+def test_that_ipython_works(ip):
+    ip.run_cell(raw_cell="x = 1").raise_error()
+    assert ip.user_global_ns["x"] == 1
+
+
+def test_function_beartype(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(x: int):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell="f(1)").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='f("x")').raise_error()
+
+
+def test_function_none(ip):
+    ip.run_cell(
+        raw_cell="""
+    def f(a,b,c):
+        pass
+                """
+    ).raise_error()
+    ip.run_cell(raw_cell='f(1,2,"k")').raise_error()
+
+
+def test_function_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_function_jaxtyped_and_jitted(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import jax
+
+    @jax.jit
+    def g(x: Float[Array, "1"]):
+        return x + 1
+
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="g(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array(1.0))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([1]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="g(jax.numpy.array([2, 3]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell='g("string")').raise_error()
+
+
+def test_class_jaxtyped(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A(eqx.Module):
+        x: Float[Array, "2"]
+
+        def do_something(self, y: Int[Array, ""]):
+            return self.x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+
+    with pytest.raises(ParamError):
+        ip.run_cell(
+            raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))"
+        ).raise_error()
+
+
+def test_class_not_dataclass(ip):
+    ip.run_cell(
+        raw_cell="""
+    from jaxtyping import Float, Array, Int
+    import equinox as eqx
+    import jax
+
+    class A:
+        def __init__(self, x):
+            self.x = x
+
+        def do_something(self, y):
+            return x + y
+                """
+    ).raise_error()
+
+    ip.run_cell(raw_cell="a = A(jax.numpy.array([1.0, 2.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array(2))").raise_error()
+    ip.run_cell(raw_cell="A(jax.numpy.array([1.0]))").raise_error()
+    ip.run_cell(raw_cell="a.do_something(jax.numpy.array([2.0, 3.0]))").raise_error()