From f454cb797c7ecf71740a19dc637b795abe75db0c Mon Sep 17 00:00:00 2001 From: Roma Knyaz Date: Wed, 20 Sep 2023 11:02:39 +0200 Subject: [PATCH] Make jaxtyping an IPython extension --- jaxtyping/__init__.py | 1 + jaxtyping/_ipython_extension.py | 35 ++++++++ test/requirements.txt | 1 + test/test_ipython_extension.py | 151 ++++++++++++++++++++++++++++++++ 4 files changed, 188 insertions(+) create mode 100644 jaxtyping/_ipython_extension.py create mode 100644 test/test_ipython_extension.py 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 ._array_types import ( ) 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." + ) 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()