Source code for k1lib.cli.nb

# AUTOGENERATED FILE! PLEASE DON'T EDIT
"""This is for everything related to ipython notebooks. Expected to use behind the
"nb" module name, like this::

    from k1lib.imports import *
    nb.execute("file.ipynb")"""
__all__ = ["cells", "pretty", "execute"]
from k1lib.cli import BaseCli; import k1lib.cli as cli
import json, k1lib, os, traceback; from typing import List
import matplotlib.pyplot as plt
def _cells(fileName, outputs=False):
    js = json.loads(cli.cat(fileName) | cli.join("\n"))
    cells = []; fields = set(["cell_type", "source"])
    if outputs: fields.add("outputs")
    return [{k:cell[k] for k in cell.keys() if k in fields} for cell in js["cells"]]
def cells(fileName=None, outputs=False):
    """Gets simplified notebook cells from file source, including fields
``cell_type`` and ``source`` only. Example::

    nb.cells("file.ipynb")"""
    if fileName is None: return cli.aS(_cells, outputs)
    else: return _cells(fileName, outputs)
class pretty(BaseCli):
    def __init__(self, magics:bool=False, whitelist:List[str]=[], blacklist:List[str]=[]):
        """Makes the cells prettier.
Cell 1 in file.ipynb::

    #notest, export
    a = 3

Cell 2 in file.ipynb::

    b = 6

Code::

    # only cell 2 gets chosen
    nb.cells("file.ipynb") | nb.pretty(blacklist=["notest"])
    # only cell 1 gets chosen
    nb.cells("file.ipynb") | nb.pretty(whitelist=["export"])

:param magics: if False, then if detected magics ('!', '%%' symbols), then remove that
    line in cell's source
:param whitelist: every cell that doesn't have any of these properties will be
    filtered out
:param blacklist: every cell that has any of these properties will be filtered out"""
        self.magics = magics; self.wl = whitelist; self.bl = blacklist
    def __ror__(self, cells):
        magics = self.magics; wl = self.wl; bl = self.bl
        for cell in cells:
            lines = cell["source"]
            if not magics: lines = [line.rstrip() for line in lines if not line.startswith("!") and not line.startswith("%%")]
            if len(lines) == 0: continue
            if lines[0].startswith("#"): props = set([e.strip() for e in lines[0].lstrip("#").split(",")])
            else: props = set()
            if len(wl) > 0 and not any([e in props for e in wl]): continue
            if len(bl) > 0 and any([e in props for e in bl]): continue
            cell["source"] = lines; yield cell
class execute(BaseCli):
    def __init__(self, fileName=None, _globals:dict=None):
        """Executes cells.
Example::

    nb.cells("file.ipynb") | nb.execute("nb.ipynb")

Most of the time, you'd want to pass cells through :class:`pretty` first, to make sure
everything is nice and clean

:param fileName: not actually used to read the file. If specified, then changes the
    current working directory to that of the file
:param _globals: optional dict of global variables"""
        self.fileName = fileName
        self._globals = dict() if _globals is None else _globals
    def __ror__(self, cells):
        with k1lib.settings.context():
            if self.fileName:
                k1lib.settings.wd = os.path.dirname(self.fileName) or "."
            for cell in cells:
                if cell["cell_type"] != "code": continue
                source = "\n".join(cell["source"] | cli.apply(lambda x: x.strip("\n")))
                try: exec(source, self._globals); plt.show()
                except Exception as e:
                    print("Problematic cell:\n")
                    print(source); traceback.print_exc(); raise e
    @staticmethod
    def rightAway(fileName:str, _globals:dict=None, tag:str=None):
        """Convenience function to execute a notebook right away.
Example::

    fn = "some/file.ipynb"
    nb.cells(fn) | nb.pretty(whitelist=[tag]) | nb.execute(nb, globals())
    nb.execute.rightAway(fn, globals(), tag)

Last 2 lines are pretty much the same."""
        if tag: tag = pretty(whitelist=[tag]) if tag else cli.iden()
        return cells(fileName) | tag | execute(fileName, _globals)