Skip to content

Commit

Permalink
Define the IPython magic lazily
Browse files Browse the repository at this point in the history
IPython is quite a chunky package, and importing it unconditionally makes
jaxtyping itself relatively slow to import.
  • Loading branch information
superbobry authored and patrick-kidger committed Apr 17, 2024
1 parent 07e58de commit f83170e
Showing 1 changed file with 11 additions and 11 deletions.
22 changes: 11 additions & 11 deletions jaxtyping/_ipython_extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,9 @@
from ._import_hook import JaxtypingTransformer, Typechecker


try:
def choose_typechecker_magics():
# The import is local to avoid degrading import times when the magic is
# not needed.
from IPython.core.magic import line_magic, Magics, magics_class

@magics_class
Expand All @@ -40,17 +42,15 @@ def typechecker(self, typechecker):
JaxtypingTransformer(typechecker=Typechecker(typechecker))
)

except Exception:
# Very broad exception-handling, as e.g. IPython will sometimes be
# present but fail to import for mysterious reasons.
pass
return ChooseTypecheckerMagics


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."
)
ChooseTypecheckerMagics = choose_typechecker_magics()
except Exception as e:
# Very broad exception-handling, as e.g. IPython will sometimes be
# present but fail to import for mysterious reasons.
raise RuntimeError("Failed to define jaxtyping.typechecker magic") from e

ipython.register_magics(ChooseTypecheckerMagics)

0 comments on commit f83170e

Please # to comment.