Source code for k1lib.cli.nb

# AUTOGENERATED FILE! PLEASE DON'T EDIT HERE. EDIT THE SOURCE NOTEBOOKS INSTEAD
"""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", "grabTags", "executeTags", "pretty", "execute"]
from k1lib.cli import BaseCli; import k1lib.cli as cli
import json, k1lib, os, traceback; from typing import List
from collections import defaultdict
plt = k1lib.dep("matplotlib.pyplot")
def _cells(fileName, outputs=False):                                             # _cells
    js = json.loads(cli.cat(fileName) | cli.join("\n"))                          # _cells
    cells = []; fields = set(["cell_type", "source"])                            # _cells
    if outputs: fields.add("outputs")                                            # _cells
    return [{k:cell[k] for k in cell.keys() if k in fields} for cell in js["cells"]] # _cells
def cells(fileName=None, outputs=False):                                         # cells
    """Gets simplified notebook cells from file source, including fields
``cell_type`` and ``source`` only. Example::

    nb.cells("file.ipynb")"""                                                    # cells
    if fileName is None: return cli.aS(_cells, outputs)                          # cells
    else: return _cells(fileName, outputs)                                       # cells
def _tagHandler(*a,**kw): pass                                                   # _tagHandler
def grabTags(s:str):                                                             # grabTags
    """Gets the tags included in a cell.
Example::

    nb.grabTags("# serve(node='mint-6.l'), abc") # returns ["serve", "abc"]
"""                                                                              # grabTags
    t = defaultdict(lambda: _tagHandler)                                         # grabTags
    s = s.strip("#").strip()                                                     # grabTags
    try: eval(f'[{s}]', t)                                                       # grabTags
    except: return [e.split("(")[0].strip() for e in s.split(",")]               # grabTags
    return [e for e in t.keys() if not (e.startswith("__") and e.endswith("__"))] # grabTags
def executeTags(s:str, funcs:dict):                                              # executeTags
    """Executes the tags included in a cell.
Example::

    d = []
    def serve(node): d.append(node)
    nb.executeTags("# serve(node='mint-6.l'), abc", {"serve": serve, "abc": 123})

    d # now is ["mint-6.l"]
"""                                                                              # executeTags
    try: s = s.strip("#").strip(); eval(f'[{s}]', dict(funcs))                   # executeTags
    except: pass                                                                 # executeTags
class pretty(BaseCli):                                                           # pretty
    def __init__(self, magics:bool=False, whitelist:List[str]=[], blacklist:List[str]=[]): # pretty
        """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""" # pretty
        self.magics = magics; self.wl = whitelist; self.bl = blacklist           # pretty
    def __ror__(self, cells):                                                    # pretty
        magics = self.magics; wl = self.wl; bl = self.bl                         # pretty
        for cell in cells:                                                       # pretty
            lines = cell["source"]                                               # pretty
            if not magics: lines = [line.rstrip() for line in lines if not line.startswith("!") and not line.startswith("%%")] # pretty
            if len(lines) == 0: continue                                         # pretty
            if lines[0].startswith("#"): props = set(grabTags(lines[0]))         # pretty
            else: props = set()                                                  # pretty
            if len(wl) > 0 and not any([e in props for e in wl]): continue       # pretty
            if len(bl) > 0 and any([e in props for e in bl]): continue           # pretty
            cell["source"] = lines; yield cell                                   # pretty
class execute(BaseCli):                                                          # execute
    def __init__(self, fileName=None, _globals:dict=None):                       # execute
        """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"""                            # execute
        self.fileName = fileName                                                 # execute
        self._globals = dict() if _globals is None else _globals                 # execute
    def __ror__(self, cells):                                                    # execute
        with k1lib.settings.context():                                           # execute
            if self.fileName:                                                    # execute
                k1lib.settings.wd = os.path.dirname(self.fileName) or "."        # execute
            for cell in cells:                                                   # execute
                if cell["cell_type"] != "code": continue                         # execute
                source = "\n".join(cell["source"] | cli.apply(lambda x: x.strip("\n"))) # execute
                try: exec(source, self._globals); plt.show()                     # execute
                except Exception as e:                                           # execute
                    print("Problematic cell:\n")                                 # execute
                    print(source); traceback.print_exc(); raise e                # execute
    @staticmethod                                                                # execute
    def rightAway(fileName:str, _globals:dict=None, tag:str="export"):           # execute
        """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."""                                        # execute
        if tag: tag = pretty(whitelist=[tag]) if tag else cli.iden()             # execute
        return cells(fileName) | tag | execute(fileName, _globals)               # execute