Skip to content

Commit

Permalink
strip ansi
Browse files Browse the repository at this point in the history
  • Loading branch information
@jph00
jph00 committed Nov 5, 2024
1 parent a21e314 commit 508b73c
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 24 deletions.
29 changes: 14 additions & 15 deletions execnb/shell.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
from IPython.core.interactiveshell import InteractiveShell, ExecutionInfo, ExecutionResult
from IPython.core.displayhook import DisplayHook
from IPython.utils.capture import capture_output
from IPython.utils.text import strip_ansi
from IPython.core.completer import IPCompleter,provisionalcompleter
from IPython.core.hooks import CommandChainDispatcher
from IPython.core.completerlib import module_completer
Expand Down Expand Up @@ -123,12 +124,12 @@ def run(self:CaptureShell,
self.exc = res.exception
return _out_nb(res, self.display_formatter)

# %% ../nbs/02_shell.ipynb 29
# %% ../nbs/02_shell.ipynb 30
def render_outputs(outputs):
def render_output(out):
otype = out['output_type']
if otype == 'stream':
txt = ''.join(out['text'])
txt = strip_ansi(''.join(out['text']))
return f"<pre>{txt}</pre>" if out['name']=='stdout' else f"<pre class='stderr'>{txt}</pre>"
elif otype in ('display_data','execute_result'):
data = out['data']
Expand All @@ -141,13 +142,11 @@ def render_output(out):
if d := _g('image/png'): return f'<img src="data:image/png;base64,{d}"/>'
if d := _g('text/latex'): return f'<div class="math">${d}$</div>'
if d := _g('text/plain'): return f"<pre>{escape(d)}</pre>"
elif otype == 'error':
return f"<pre class='error'>{out['ename']}: {out['evalue']}\n{''.join(out['traceback'])}</pre>"
return ''

return '\n'.join(map(render_output, outputs))

# %% ../nbs/02_shell.ipynb 40
# %% ../nbs/02_shell.ipynb 41
@patch
def cell(self:CaptureShell, cell, stdout=True, stderr=True):
"Run `cell`, skipping if not code, and store outputs back in cell"
Expand All @@ -159,32 +158,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 43
# %% ../nbs/02_shell.ipynb 44
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 46
# %% ../nbs/02_shell.ipynb 47
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 48
# %% ../nbs/02_shell.ipynb 49
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 50
# %% ../nbs/02_shell.ipynb 51
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 52
# %% ../nbs/02_shell.ipynb 53
def _false(o): return False

@patch
Expand All @@ -204,7 +203,7 @@ def run_all(self:CaptureShell,
postproc(cell)
if self.exc and exc_stop: raise self.exc from None

# %% ../nbs/02_shell.ipynb 66
# %% ../nbs/02_shell.ipynb 67
@patch
def execute(self:CaptureShell,
src:str|Path, # Notebook path to read from
Expand All @@ -225,7 +224,7 @@ def execute(self:CaptureShell,
inject_code=inject_code, inject_idx=inject_idx)
if dest: write_nb(nb, dest)

# %% ../nbs/02_shell.ipynb 70
# %% ../nbs/02_shell.ipynb 71
@patch
def prettytb(self:CaptureShell,
fname:str|Path=None): # filename to print alongside the traceback
Expand All @@ -237,7 +236,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 89
# %% ../nbs/02_shell.ipynb 90
@call_parse
def exec_nb(
src:str, # Notebook path to read from
Expand All @@ -251,7 +250,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 92
# %% ../nbs/02_shell.ipynb 93
class SmartCompleter(IPCompleter):
def __init__(self, shell, namespace=None, jedi=False):
if namespace is None: namespace = shell.user_ns
Expand All @@ -271,7 +270,7 @@ def __call__(self, c):
for o in self.completions(c, len(c))
if o.type=='<unknown>']

# %% ../nbs/02_shell.ipynb 94
# %% ../nbs/02_shell.ipynb 95
@patch
def complete(self:CaptureShell, c):
if not hasattr(self, '_completer'): self._completer = SmartCompleter(self)
Expand Down
52 changes: 43 additions & 9 deletions nbs/02_shell.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@
"from IPython.core.interactiveshell import InteractiveShell, ExecutionInfo, ExecutionResult\n",
"from IPython.core.displayhook import DisplayHook\n",
"from IPython.utils.capture import capture_output\n",
"from IPython.utils.text import strip_ansi\n",
"from IPython.core.completer import IPCompleter,provisionalcompleter\n",
"from IPython.core.hooks import CommandChainDispatcher\n",
"from IPython.core.completerlib import module_completer\n",
Expand Down Expand Up @@ -706,8 +707,8 @@
"text/plain": [
"[{'name': 'stdout',\n",
" 'output_type': 'stream',\n",
" 'text': ['CPU times: user 1e+03 ns, sys: 1e+03 ns, total: 2 us\\n',\n",
" 'Wall time: 3.1 us\\n']},\n",
" 'text': ['CPU times: user 2 us, sys: 0 ns, total: 2 us\\n',\n",
" 'Wall time: 3.81 us\\n']},\n",
" {'data': {'text/plain': ['2']},\n",
" 'metadata': {},\n",
" 'output_type': 'execute_result'}]"
Expand All @@ -723,6 +724,38 @@
"o"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"[{'name': 'stdout',\n",
" 'output_type': 'stream',\n",
" 'text': ['\\x1b[0;31m---------------------------------------------------------------------------\\x1b[0m\\n',\n",
" '\\x1b[0;31mZeroDivisionError\\x1b[0m Traceback (most recent call last)\\n',\n",
" 'File \\x1b[0;32m<ipython-input-1-9e1622b385b6>:1\\x1b[0m\\n',\n",
" '\\x1b[0;32m----> 1\\x1b[0m \\x1b[38;5;241m1\\x1b[39m\\x1b[38;5;241m/\\x1b[39m\\x1b[38;5;241m0\\x1b[39m\\n',\n",
" '\\n',\n",
" '\\x1b[0;31mZeroDivisionError\\x1b[0m: division by zero\\n']},\n",
" {'ename': 'ZeroDivisionError',\n",
" 'evalue': 'division by zero',\n",
" 'output_type': 'error',\n",
" 'traceback': 'Traceback (most recent call last):\\n File \"/Users/jhoward/miniconda3/lib/python3.11/site-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\\n exec(code_obj, self.user_global_ns, self.user_ns)\\n File \"<ipython-input-1-9e1622b385b6>\", line 1, in <module>\\n 1/0\\n ~^~\\nZeroDivisionError: division by zero\\n'}]"
]
},
"execution_count": null,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"o = s.run(\"1/0\")\n",
"o"
]
},
{
"cell_type": "code",
"execution_count": null,
Expand All @@ -734,7 +767,7 @@
" def render_output(out):\n",
" otype = out['output_type']\n",
" if otype == 'stream':\n",
" txt = ''.join(out['text'])\n",
" txt = strip_ansi(''.join(out['text']))\n",
" return f\"<pre>{txt}</pre>\" if out['name']=='stdout' else f\"<pre class='stderr'>{txt}</pre>\"\n",
" elif otype in ('display_data','execute_result'):\n",
" data = out['data']\n",
Expand All @@ -747,8 +780,6 @@
" if d := _g('image/png'): return f'<img src=\"data:image/png;base64,{d}\"/>'\n",
" if d := _g('text/latex'): return f'<div class=\"math\">${d}$</div>'\n",
" if d := _g('text/plain'): return f\"<pre>{escape(d)}</pre>\"\n",
" elif otype == 'error':\n",
" return f\"<pre class='error'>{out['ename']}: {out['evalue']}\\n{''.join(out['traceback'])}</pre>\"\n",
" return ''\n",
" \n",
" return '\\n'.join(map(render_output, outputs))"
Expand All @@ -762,10 +793,13 @@
{
"data": {
"text/html": [
"<pre>CPU times: user 1e+03 ns, sys: 1e+03 ns, total: 2 us\n",
"Wall time: 3.1 us\n",
"</pre>\n",
"<pre>2</pre>"
"<pre>---------------------------------------------------------------------------\n",
"ZeroDivisionError Traceback (most recent call last)\n",
"File <ipython-input-1-9e1622b385b6>:1\n",
"----> 1 1/0\n",
"\n",
"ZeroDivisionError: division by zero\n",
"</pre>\n"
],
"text/plain": [
"<IPython.core.display.HTML object>"
Expand Down

0 comments on commit 508b73c

Please sign in to comment.