fixes #60

Author: jph00

execnb/nbio.py

@@ -5,19 +5,19 @@
 # %% auto 0
 __all__ = ['NbCell', 'dict2nb', 'read_nb', 'new_nb', 'mk_cell', 'nb2dict', 'nb2str', 'write_nb']
 
-# %% ../nbs/01_nbio.ipynb 2
+# %% ../nbs/01_nbio.ipynb
 from fastcore.basics import *
 from fastcore.imports import *
 
 import ast,functools
 from pprint import pformat,pprint
 from json import loads,dumps
 
-# %% ../nbs/01_nbio.ipynb 6
+# %% ../nbs/01_nbio.ipynb
 def _read_json(self, encoding=None, errors=None):
     return loads(Path(self).read_text(encoding=encoding, errors=errors))
 
-# %% ../nbs/01_nbio.ipynb 13
+# %% ../nbs/01_nbio.ipynb
 class NbCell(AttrDict):
     def __init__(self, idx, cell):
         super().__init__(cell)
@@ -40,7 +40,7 @@ def parsed_(self):
     def __hash__(self): return hash(self.source) + hash(self.cell_type)
     def __eq__(self,o): return self.source==o.source and self.cell_type==o.cell_type
 
-# %% ../nbs/01_nbio.ipynb 15
+# %% ../nbs/01_nbio.ipynb
 def _dict2obj(d, list_func=list, dict_func=AttrDict):
     "Convert (possibly nested) dicts (or lists of dicts) to `AttrDict`"
     if isinstance(d, list): return list(map(_dict2obj, d))
@@ -53,19 +53,19 @@ def dict2nb(js=None, **kwargs):
     nb.cells = [NbCell(*o) for o in enumerate(nb.cells)]
     return nb
 
-# %% ../nbs/01_nbio.ipynb 20
+# %% ../nbs/01_nbio.ipynb
 def read_nb(path):
     "Return notebook at `path`"
     res = dict2nb(_read_json(path, encoding='utf-8'))
     res['path_'] = str(path)
     return res
 
-# %% ../nbs/01_nbio.ipynb 26
+# %% ../nbs/01_nbio.ipynb
 def new_nb(cells=None, meta=None, nbformat=4, nbformat_minor=5):
     "Returns an empty new notebook"
     return dict2nb(cells=cells or [],metadata=meta or {},nbformat=nbformat,nbformat_minor=nbformat_minor)
 
-# %% ../nbs/01_nbio.ipynb 28
+# %% ../nbs/01_nbio.ipynb
 def mk_cell(text,  # `source` attr in cell
             cell_type='code',  # `cell_type` attr in cell
             **kwargs):  # any other attrs to add to cell
@@ -77,21 +77,21 @@ def mk_cell(text,  # `source` attr in cell
         kwargs['execution_count']=0
     return NbCell(0, dict(cell_type=cell_type, source=text, directives_={}, **kwargs))
 
-# %% ../nbs/01_nbio.ipynb 31
+# %% ../nbs/01_nbio.ipynb
 def nb2dict(d, k=None):
     "Convert parsed notebook to `dict`"
     if k=='source': return d.splitlines(keepends=True)
     if isinstance(d, list): return list(map(nb2dict,d))
     if not isinstance(d, dict): return d
     return dict(**{k:nb2dict(v,k) for k,v in d.items() if k[-1] != '_'})
 
-# %% ../nbs/01_nbio.ipynb 34
+# %% ../nbs/01_nbio.ipynb
 def nb2str(nb):
     "Convert `nb` to a `str`"
     if isinstance(nb, (AttrDict,list)): nb = nb2dict(nb)
     return dumps(nb, sort_keys=True, indent=1, ensure_ascii=False) + "\n"
 
-# %% ../nbs/01_nbio.ipynb 37
+# %% ../nbs/01_nbio.ipynb
 def write_nb(nb, path):
     "Write `nb` to `path`"
     new = nb2str(nb)

execnb/shell.py

@@ -9,7 +9,7 @@
 from fastcore.script import call_parse
 from fastcore.ansi import ansi2html
 
-import mistletoe,multiprocessing,types,traceback
+import mistletoe,multiprocessing,types,traceback,signal
 try:
     if sys.platform == 'darwin': multiprocessing.set_start_method("fork")
 except RuntimeError: pass # if re-running cell
@@ -37,7 +37,7 @@
 __all__ = ['CaptureShell', 'format_exc', 'render_outputs', 'find_output', 'out_exec', 'out_stream', 'out_error', 'exec_nb',
            'SmartCompleter']
 
-# %% ../nbs/02_shell.ipynb 5
+# %% ../nbs/02_shell.ipynb
 class _CustDisplayHook(DisplayHook):
     def write_output_prompt(self): pass
     def write_format_data(self, data, md_dict): pass
@@ -49,12 +49,14 @@ def __repr__(self: ExecutionInfo): return f'cell: {self.raw_cell}; id: {self.cel
 @patch
 def __repr__(self: ExecutionResult): return f'result: {self.result}; err: {self.error_in_exec}; info: <{self.info}>'
 
-# %% ../nbs/02_shell.ipynb 6
+# %% ../nbs/02_shell.ipynb
 class CaptureShell(InteractiveShell):
     displayhook_class = _CustDisplayHook
 
-    def __init__(self, path:str|Path=None, mpl_format='retina'):
+    def __init__(self, path:str|Path=None, mpl_format='retina', history=False, timeout=None):
         super().__init__()
+        self.history_manager.enabled = history
+        self.timeout = timeout
         self.result,self.exc = None,None
         if path: self.set_path(path)
         self.display_formatter.active = True
@@ -65,11 +67,17 @@ def __init__(self, path:str|Path=None, mpl_format='retina'):
             self.run_cell(f"set_matplotlib_formats('{mpl_format}')")
 
     def run_cell(self, raw_cell, store_history=False, silent=False, shell_futures=True, cell_id=None,
-                 stdout=True, stderr=True, display=True):
+                 stdout=True, stderr=True, display=True, timeout=None):
+        if not timeout: timeout = self.timeout
         # TODO what if there's a comment?
         semic = raw_cell.rstrip().endswith(';')
+        if timeout:
+            def handler(*args): raise TimeoutError()
+            signal.signal(signal.SIGALRM, handler)
+            signal.alarm(timeout)
         with capture_output(display=display, stdout=stdout, stderr=stderr) as c:
             result = super().run_cell(raw_cell, store_history, silent, shell_futures=shell_futures, cell_id=cell_id)
+        if timeout: signal.alarm(0)
         return AttrDict(result=result, stdout='' if semic else c.stdout, stderr=c.stderr,
                         display_objects=c.outputs, exception=result.error_in_exec, quiet=semic)
 
@@ -81,12 +89,12 @@ def set_path(self, path):
 
     def enable_gui(self, gui=None): pass
 
-# %% ../nbs/02_shell.ipynb 23
+# %% ../nbs/02_shell.ipynb
 def format_exc(e):
     "Format exception `e` as a string"
     return ''.join(traceback.format_exception(type(e), e, e.__traceback__))
 
-# %% ../nbs/02_shell.ipynb 24
+# %% ../nbs/02_shell.ipynb
 def _out_stream(text, name): return dict(name=name, output_type='stream', text=text.splitlines(True))
 def _out_exc(e):
     ename = type(e).__name__
@@ -116,7 +124,7 @@ def _out_nb(o, fmt):
         res.append(_mk_out(*fmt.format(r), 'execute_result'))
     return res
 
-# %% ../nbs/02_shell.ipynb 25
+# %% ../nbs/02_shell.ipynb
 @patch
 def run(self:CaptureShell,
         code:str, # Python/IPython code to run
@@ -128,7 +136,7 @@ def run(self:CaptureShell,
     self.exc = res.exception
     return _out_nb(res, self.display_formatter)
 
-# %% ../nbs/02_shell.ipynb 31
+# %% ../nbs/02_shell.ipynb
 def render_outputs(outputs, ansi_renderer=strip_ansi):
     def render_output(out):
         otype = out['output_type']
@@ -150,7 +158,7 @@ def render_output(out):
 
     return '\n'.join(map(render_output, outputs))
 
-# %% ../nbs/02_shell.ipynb 44
+# %% ../nbs/02_shell.ipynb
 @patch
 def cell(self:CaptureShell, cell, stdout=True, stderr=True):
     "Run `cell`, skipping if not code, and store outputs back in cell"
@@ -162,32 +170,32 @@ def cell(self:CaptureShell, cell, stdout=True, stderr=True):
         for o in outs:
             if 'execution_count' in o: cell['execution_count'] = o['execution_count']
 
-# %% ../nbs/02_shell.ipynb 47
+# %% ../nbs/02_shell.ipynb
 def find_output(outp, # Output from `run`
                 ot='execute_result' # Output_type to find
                ):
     "Find first output of type `ot` in `CaptureShell.run` output"
     return first(o for o in outp if o['output_type']==ot)
 
-# %% ../nbs/02_shell.ipynb 50
+# %% ../nbs/02_shell.ipynb
 def out_exec(outp):
     "Get data from execution result in `outp`."
     out = find_output(outp)
     if out: return '\n'.join(first(out['data'].values()))
 
-# %% ../nbs/02_shell.ipynb 52
+# %% ../nbs/02_shell.ipynb
 def out_stream(outp):
     "Get text from stream in `outp`."
     out = find_output(outp, 'stream')
     if out: return ('\n'.join(out['text'])).strip()
 
-# %% ../nbs/02_shell.ipynb 54
+# %% ../nbs/02_shell.ipynb
 def out_error(outp):
     "Get traceback from error in `outp`."
     out = find_output(outp, 'error')
     if out: return '\n'.join(out['traceback'])
 
-# %% ../nbs/02_shell.ipynb 56
+# %% ../nbs/02_shell.ipynb
 def _false(o): return False
 
 @patch
@@ -207,7 +215,7 @@ def run_all(self:CaptureShell,
             postproc(cell)
         if self.exc and exc_stop: raise self.exc from None
 
-# %% ../nbs/02_shell.ipynb 70
+# %% ../nbs/02_shell.ipynb
 @patch
 def execute(self:CaptureShell,
             src:str|Path, # Notebook path to read from
@@ -228,7 +236,7 @@ def execute(self:CaptureShell,
                  inject_code=inject_code, inject_idx=inject_idx)
     if dest: write_nb(nb, dest)
 
-# %% ../nbs/02_shell.ipynb 74
+# %% ../nbs/02_shell.ipynb
 @patch
 def prettytb(self:CaptureShell, 
              fname:str|Path=None): # filename to print alongside the traceback
@@ -240,7 +248,7 @@ def prettytb(self:CaptureShell,
     fname_str = f' in {fname}' if fname else ''
     return f"{type(self.exc).__name__}{fname_str}:\n{_fence}\n{cell_str}\n"
 
-# %% ../nbs/02_shell.ipynb 93
+# %% ../nbs/02_shell.ipynb
 @call_parse
 def exec_nb(
     src:str, # Notebook path to read from
@@ -254,7 +262,7 @@ def exec_nb(
     CaptureShell().execute(src, dest, exc_stop=exc_stop, inject_code=inject_code,
                            inject_path=inject_path, inject_idx=inject_idx)
 
-# %% ../nbs/02_shell.ipynb 96
+# %% ../nbs/02_shell.ipynb
 class SmartCompleter(IPCompleter):
     def __init__(self, shell, namespace=None, jedi=False):
         if namespace is None: namespace = shell.user_ns
@@ -274,7 +282,7 @@ def __call__(self, c):
                     for o in self.completions(c, len(c))
                     if o.type not in ('magic', 'path')]
 
-# %% ../nbs/02_shell.ipynb 98
+# %% ../nbs/02_shell.ipynb
 @patch
 def complete(self:CaptureShell, c):
     if not hasattr(self, '_completer'): self._completer = SmartCompleter(self)

nbs/02_shell.ipynb

@@ -39,4 +39,5 @@ jupyter_hooks = True
 clean_ids = True
 clear_all = False
 put_version_in_init = True
+cell_number = False