"""
Utilities for handling Jupyter / IPython notebooks
This code is copied and modified from nbimporter
(https://github.com/grst/nbimporter/blob/master/nbimporter.py) which is not
actively maintained (otherwise we would use it as a dependency).
Note that using this behavior is very much discouraged, it would be far better
if you maintained your reusable code in separate python modules. See
https://github.com/grst/nbimporter for reasons.
----
Allow for importing of IPython Notebooks as modules from Jupyter v4.
Updated from module collated here:
https://github.com/adrn/ipython/blob/master/examples/Notebook/Importing%20Notebooks.ipynb
Importing from a notebook is different from a module: because one
typically keeps many computations and tests besides exportable defs,
here we only run code which either defines a function or a class, or
imports code from other modules and notebooks. This behaviour can be
disabled by setting NotebookLoader.default_options['only_defs'] = False.
Furthermore, in order to provide per-notebook initialisation, if a
special function __nbinit__() is defined in the notebook, it will be
executed the first time an import statement is. This behaviour can be
disabled by setting NotebookLoader.default_options['run_nbinit'] = False.
Finally, you can set the encoding of the notebooks with
NotebookLoader.default_options['encoding']. The default is 'utf-8'.
"""
from __future__ import annotations
import ast
import io
import os
import sys
import types
import typing
from os.path import basename, dirname
[docs]
def _find_notebook(fullname, path=None):
"""Find a notebook, given its fully qualified name and an optional path
This turns "foo.bar" into "foo/bar.ipynb"
and tries turning "Foo_Bar" into "Foo Bar" if Foo_Bar
does not exist.
Example:
>>> # xdoctest: +REQUIRES(PY3, module:IPython, module:nbconvert)
>>> from xdoctest.utils.util_notebook import _find_notebook
>>> from xdoctest import utils
>>> from os.path import join, basename, splitext
>>> self = utils.TempDir()
>>> dpath = self.ensure()
>>> fpath = join(dpath, 'test_import_notebook.ipynb')
>>> cells = ['x = 1']
>>> _make_test_notebook_fpath(fpath, cells)
>>> fullname = splitext(basename(fpath))[0]
>>> path = [dpath]
>>> _find_notebook(fullname, path)
...test_import_notebook.ipynb
>>> _find_notebook(fullname, None)
None
"""
name = fullname.rsplit('.', 1)[-1]
if not path:
path = ['']
for d in path:
nb_path = os.path.join(d, name + '.ipynb')
if os.path.isfile(nb_path):
return nb_path
# let import Notebook_Name find "Notebook Name.ipynb"
nb_path = nb_path.replace('_', ' ')
if os.path.isfile(nb_path):
return nb_path
[docs]
class CellDeleter(ast.NodeTransformer):
"""Removes all nodes from an AST which are not suitable
for exporting out of a notebook."""
[docs]
def visit(self, node):
"""Visit a node."""
if node.__class__.__name__ in [
'Module',
'FunctionDef',
'ClassDef',
'Import',
'ImportFrom',
]:
return node
return None
[docs]
class NotebookLoader:
"""Module Loader for Jupyter Notebooks."""
default_options = {
'only_defs': False,
'run_nbinit': True,
'encoding': 'utf-8',
}
def __init__(self, path=None) -> None:
from IPython.core.interactiveshell import InteractiveShell
self.shell = InteractiveShell.instance()
self.path = path
self.options = self.default_options.copy()
[docs]
def load_module(
self,
fullname: str | None = None,
fpath: str | os.PathLike | None = None,
):
"""import a notebook as a module"""
import nbformat
from IPython import get_ipython
if fpath is None:
fpath = _find_notebook(fullname, self.path)
# load the notebook object
nb_version = nbformat.current_nbformat
assert fpath is not None
fpath_str = str(fpath)
with io.open(fpath_str, 'r', encoding=self.options['encoding']) as f: # type: ignore
nb = nbformat.read(f, nb_version)
# create the module and add it to sys.modules
# if name in sys.modules:
# return sys.modules[name]
assert isinstance(fullname, str)
mod = types.ModuleType(fullname)
mod.__file__ = fpath_str
mod.__loader__ = self
mod.__dict__['get_ipython'] = get_ipython
# Only do something if it's a python notebook
# if nb.metadata.kernelspec.language != 'python':
# print("Ignoring '%s': not a python notebook." % fpath)
# return mod
# print("Importing Jupyter notebook from %s" % fpath)
sys.modules[fullname] = mod
# extra work to ensure that magics that would affect the user_ns
# actually affect the notebook module's ns
save_user_ns = self.shell.user_ns
self.shell.user_ns = mod.__dict__
try:
deleter = CellDeleter()
for cell in filter(lambda c: c.cell_type == 'code', nb.cells):
# transform the input into executable Python
code = self.shell.input_transformer_manager.transform_cell(
cell.source
)
tree: ast.AST
if self.options['only_defs']:
# Remove anything that isn't a def or a class
tree = deleter.generic_visit(ast.parse(code))
else:
tree = typing.cast(ast.AST, ast.parse(code))
# run the code in the module
codeobj = compile(
typing.cast(
typing.Union[
ast.Module, ast.Expression, ast.Interactive
],
tree,
),
filename=fpath,
mode='exec',
)
exec(codeobj, mod.__dict__)
finally:
self.shell.user_ns = save_user_ns
# Run any initialisation if available, but only once
if self.options['run_nbinit'] and '__nbinit_done__' not in mod.__dict__:
try:
mod.__nbinit__()
setattr(mod, '__nbinit_done__', True)
except (KeyError, AttributeError):
pass
return mod
[docs]
def import_notebook_from_path(
ipynb_fpath: str | os.PathLike, only_defs: bool = False
):
"""
Import an IPython notebook as a module from a full path and try to maintain
clean sys.path variables.
Args:
ipynb_fpath (str | PathLike): path to the ipython notebook file to import
only_defs (bool, default=False): if True ignores all non-definition
statements
Example:
>>> # xdoctest: +REQUIRES(PY3, module:IPython, module:nbconvert)
>>> from xdoctest import utils
>>> from os.path import join
>>> self = utils.TempDir()
>>> dpath = self.ensure()
>>> ipynb_fpath = join(dpath, 'test_import_notebook.ipydb')
>>> cells = [
>>> utils.codeblock(
>>> '''
>>> def foo():
>>> return 'bar'
>>> '''),
>>> utils.codeblock(
>>> '''
>>> x = 1
>>> ''')
>>> ]
>>> _make_test_notebook_fpath(ipynb_fpath, cells)
>>> module = import_notebook_from_path(ipynb_fpath)
>>> assert module.foo() == 'bar'
>>> assert module.x == 1
"""
ipynb_fname = basename(ipynb_fpath)
fname_noext = ipynb_fname.rsplit('.', 1)[0]
ipynb_modname = fname_noext.replace(' ', '_')
# hack around the importlib machinery
loader = NotebookLoader()
loader.options['only_defs'] = only_defs
module = loader.load_module(ipynb_modname, ipynb_fpath)
return module
[docs]
def execute_notebook(
ipynb_fpath: str | os.PathLike,
timeout: typing.Any = None,
verbose: bool | int | None = None,
) -> tuple[typing.Any, dict[str, object]]:
"""
Execute an IPython notebook in a separate kernel
Args:
ipynb_fpath (str | PathLike): path to the ipython notebook file to import
Returns:
nbformat.notebooknode.NotebookNode : nb
The executed notebook.
dict: resources
Additional resources used in the conversion process.
Example:
>>> # xdoctest: +REQUIRES(PY3, module:IPython, module:nbconvert, CPYTHON)
>>> from xdoctest import utils
>>> from os.path import join
>>> self = utils.TempDir()
>>> dpath = self.ensure()
>>> ipynb_fpath = join(dpath, 'hello_world.ipydb')
>>> _make_test_notebook_fpath(ipynb_fpath, [utils.codeblock(
>>> '''
>>> print('hello world')
>>> ''')])
>>> nb, resources = execute_notebook(ipynb_fpath, verbose=3)
>>> print('resources = {!r}'.format(resources))
>>> print('nb = {!r}'.format(nb))
>>> for cell in nb['cells']:
>>> if len(cell['outputs']) != 1:
>>> import warnings
>>> warnings.warn('expected an output, is this the issue '
>>> 'described [here](https://github.com/nteract/papermill/issues/426)?')
"""
import logging
import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
dpath = dirname(ipynb_fpath)
ep = ExecutePreprocessor(timeout=timeout)
if verbose is None:
verbose = 0
if verbose > 1:
print('executing notebook in dpath = {!r}'.format(dpath))
ep.log.setLevel(logging.DEBUG)
elif verbose > 0:
ep.log.setLevel(logging.INFO)
with open(ipynb_fpath, 'r+') as file:
nb = nbformat.read(file, as_version=nbformat.NO_CONVERT)
nb, resources = ep.preprocess(nb, {'metadata': {'path': dpath}})
# from nbconvert.preprocessors import executenb
# nb, resources = executenb(nb, cwd=dpath)
return nb, resources
[docs]
def _make_test_notebook_fpath(
fpath: typing.Any, cell_sources: typing.Any
) -> typing.Any:
"""
Helper for testing
Args:
fpath (str): file to write notebook to
cell_sources (List[str]): list of python code blocks
References:
https://stackoverflow.com/questions/38193878/create-notebook-from-code
https://gist.github.com/fperez/9716279
"""
import json
import jupyter_client.kernelspec
import nbformat as nbf
# TODO: is there an API to generate kernelspec json correctly?
kernel_name = jupyter_client.kernelspec.NATIVE_KERNEL_NAME
spec = jupyter_client.kernelspec.get_kernel_spec(kernel_name)
metadata = {
'kernelspec': {
'name': kernel_name,
'display_name': spec.display_name,
'language': spec.language,
}
}
# Use nbformat API to create notebook structure and cell json
nb = nbf.v4.new_notebook(metadata=metadata)
for source in cell_sources:
nb['cells'].append(nbf.v4.new_code_cell(source))
with open(fpath, 'w') as file:
json.dump(nb, file)
return fpath
if __name__ == '__main__':
"""
CommandLine:
python ~/code/xdoctest/xdoctest/utils/util_notebook.py all
"""
import xdoctest
xdoctest.doctest_module(__file__)