def test_wrap_code_cells_2():
# Test wrap_code_cells() with two code cells.
ast = ASTNode('root')
cb0 = ASTNode('CodeBlock', lang='python', children=['a'])
cb1 = ASTNode('CodeBlock', lang='python', children=['b'])
ast.add_child(cb0)
ast.add_child(cb1)
ast.show()
ast_wrapped = wrap_code_cells(ast)
ast_wrapped.show()
ast_expected = ASTNode('root')
# First code cell.
code_cell0 = ASTNode('CodeCell')
code_cell0.add_child(cb0)
ast_expected.add_child(code_cell0)
# Second code cell.
code_cell1 = ASTNode('CodeCell')
code_cell1.add_child(cb1)
ast_expected.add_child(code_cell1)
ast_expected.show()
assert ast_wrapped == ast_expected
def test_wrap_code_cells_2():
# Test wrap_code_cells() with two code cells.
ast = ASTNode('root')
cb0 = ASTNode('CodeBlock', lang='python', children=['a'])
cb1 = ASTNode('CodeBlock', lang='python', children=['b'])
ast.add_child(cb0)
ast.add_child(cb1)
ast.show()
ast_wrapped = wrap_code_cells(ast)
ast_wrapped.show()
ast_expected = ASTNode('root')
# First code cell.
code_cell0 = ASTNode('CodeCell')
code_cell0.add_child(cb0)
ast_expected.add_child(code_cell0)
# Second code cell.
code_cell1 = ASTNode('CodeCell')
code_cell1.add_child(cb1)
ast_expected.add_child(code_cell1)
ast_expected.show()
assert_equal(ast_wrapped, ast_expected)
def read_code(self, cell, cell_index=None):
node = ASTNode('CodeCell')
# The first child is the source.
# NOTE: the language of the code block is the notebook's language.
node.add_child(ASTNode('CodeBlock',
lang=self.language,
children=[cell.source]))
# Then, we add one extra child per output.
for output_index, output in enumerate(cell.get('outputs', [])):
if output.output_type == 'stream':
child = ASTNode('CodeBlock', lang=output.name, # stdout/stderr
children=[output.text])
elif output.output_type in ('display_data', 'execute_result'):
# Output text node.
text = output.data.get('text/plain', 'Output')
# Extract image output, if any.
out = extract_output(output)
if out is None:
child = ASTNode('CodeBlock', lang='result',
children=[text])
else:
mime_type, data = out
fn = output_filename(mime_type=mime_type,
cell_index=cell_index,
output_index=output_index,
unique_key=None, # TODO
)
self.resources[fn] = data
child = ASTNode('Image', url=fn, children=[text])
node.add_child(child)
self.tree.children.append(node)
def test_notebook_empty():
ast = ASTNode('root')
ast.add_child(ASTNode('Para'))
nb = NotebookWriter().write(ast)
reader = NotebookReader()
ast = reader.read(nb)
assert not ast.children
class CodeCellWrapper(object):
def infer_language(self, ast):
"""Return the most common CodeBlock language: it is supposed to be the
notebook's language."""
mc = Counter([_get_cell_lang(node) for node in ast.children
if node.name in ('CodeBlock', 'CodeCell')]).most_common(1)
return mc[0][0] if mc else 'python'
def wrap(self, ast):
self.ast = ast.copy()
self.ast.children = []
self._code_cell = None
# Infer the notebook's language.
self.language = self.infer_language(ast)
for i, node in enumerate(ast.children):
if self._code_cell:
if self.is_output(node) or self.is_image(node):
self.add_output(node)
else:
self.end_code_cell()
if not self._code_cell:
if self.is_source(node):
self.start_code_cell(node)
else:
self.append(node)
# Ensure the last cell is appended.
self.end_code_cell()
return self.ast
def is_output(self, node):
return ((node.name == 'CodeBlock') and
(node.lang in (None, '') or node.lang.startswith('{output')))
def is_image(self, node):
children = node.children
return ((node.name == 'Para') and
(len(children) == 1) and
(isinstance(children[0], ASTNode)) and
(children[0].name == 'Image'))
def is_source(self, node):
return node.name == 'CodeBlock' and node.lang == self.language
def start_code_cell(self, node):
self._code_cell = ASTNode('CodeCell')
# Source CodeBlock.
self._code_cell.add_child(node)
def add_output(self, node):
self._code_cell.add_child(node)
def end_code_cell(self):
if self._code_cell:
self.append(self._code_cell)
self._code_cell = None
def append(self, node):
self.ast.add_child(node)
def test_wrap_code_cells_1():
# Test wrap_code_cells() with a single code cell.
ast = ASTNode('root')
ast.add_child(ASTNode('CodeBlock', lang='python', children=['']))
ast_wrapped = wrap_code_cells(ast)
ast_wrapped.show()
ast_expected = ASTNode('root')
ast_expected.add_child(ASTNode('CodeCell', children=[ast.children[0]]))
assert_equal(ast_wrapped, ast_expected)
def test_wrap_code_cells_1():
# Test wrap_code_cells() with a single code cell.
ast = ASTNode('root')
ast.add_child(ASTNode('CodeBlock', lang='python', children=['']))
ast_wrapped = wrap_code_cells(ast)
ast_wrapped.show()
ast_expected = ASTNode('root')
ast_expected.add_child(ASTNode('CodeCell', children=[ast.children[0]]))
assert ast_wrapped == ast_expected
def read_code(self, cell, cell_index=None):
node = ASTNode('CodeCell')
# TODO: improve this.
node._visit_meta['is_block'] = True
# The first child is the source.
# NOTE: the language of the code block is the notebook's language.
node.add_child(ASTNode('CodeBlock',
lang=self.language,
children=[cell.source.rstrip()]))
# Then, we add one extra child per output.
for output_index, output in enumerate(cell.get('outputs', [])):
if output.output_type == 'stream':
child = ASTNode('CodeBlock',
lang='{output:' + output.name + '}', # stdout/stderr
children=[_remove_ansi(output.text.rstrip())])
elif output.output_type == 'error': # pragma: no cover
child = ASTNode('CodeBlock',
lang='{output:error}',
children=[_remove_ansi('\n'.join(output.traceback))])
elif output.output_type in ('display_data', 'execute_result'):
# Output text node.
# Take it from cell metadata first, otherwise from the cell's output text.
text = cell.metadata.get('podoc', {}).get('output_text', None)
text = text or output.data.get('text/plain', 'Output')
# Remove color codes.
text = _remove_ansi(text)
# Extract image output, if any.
out = extract_image(output)
if out is None:
out = extract_table(output)
if out is None:
child = ASTNode('CodeBlock',
lang='{output:result}',
children=[text])
else:
mime_type, data = out
fn = output_filename(mime_type=mime_type,
cell_index=cell_index,
output_index=output_index,
unique_key=self._unique_key,
)
self.resources[fn] = data
# Prevent multi-line image legend.
if '\n' in text: # pragma: no cover
text = 'Output'
# Wrap the Image node in a Para.
img_child = ASTNode('Image', url='{resource:%s}' % fn,
children=[text])
child = ASTNode('Para', children=[img_child])
else: # pragma: no cover
raise ValueError("Unknown output type `%s`." % output.output_type)
node.add_child(child)
self.tree.children.append(node)
class CodeCellWrapper(object):
def wrap(self, ast):
self.ast = ast.copy()
self.ast.children = []
self._code_cell = None
for i, node in enumerate(ast.children):
if self._code_cell:
if self.is_output(node) or self.is_image(node):
self.add_output(node)
else:
self.end_code_cell()
if not self._code_cell:
if self.is_source(node):
self.start_code_cell(node)
else:
self.append(node)
# Ensure the last cell is appended.
self.end_code_cell()
return self.ast
def is_output(self, node):
return ((node.name == 'CodeBlock') and
(node.lang in (None, '', 'stdout', 'stderr', 'result')))
def is_image(self, node):
children = node.children
return ((node.name == 'Para') and
(len(children) == 1) and
(isinstance(children[0], ASTNode)) and
(children[0].name == 'Image'))
def is_source(self, node):
# TODO: customizable lang
return node.name == 'CodeBlock' and node.lang == 'python'
def start_code_cell(self, node):
self._code_cell = ASTNode('CodeCell')
# Source CodeBlock.
self._code_cell.add_child(node)
def add_output(self, node):
self._code_cell.add_child(node)
def end_code_cell(self):
if self._code_cell:
self.append(self._code_cell)
self._code_cell = None
def append(self, node):
self.ast.add_child(node)
def wrap_code_cells(ast):
"""Take an AST and wrap top-level CodeBlocks within CodeCells."""
out = ast.copy()
out.children = []
current_cell = None
for i, child in enumerate(ast.children):
# Notebook code cell.
if child.name == 'CodeBlock' and child.lang == 'python':
current_cell = ASTNode('CodeCell')
# TODO: parameterizable language
# Wrap CodeBlocks within CodeCells.
current_cell.add_child(child)
else:
# Decide whether we're part of the current cell.
name = child.name
children = child.children
# Case 1: we're a code block with a notebook-specific language.
is_output = ((name == 'CodeBlock') and
(child.lang in (None, '', 'stdout',
'stderr', 'result')))
# Case 2: we're just an image.
is_image = ((name == 'Para') and
(len(children) == 1) and
(isinstance(children[0], ASTNode)) and
(children[0].name == 'Image'))
if current_cell:
if is_output or is_image:
# Add the current block to the cell's outputs.
current_cell.add_child(child)
else:
# We're no longer part of the current cell.
# First, we add the cell that has just finished.
out.add_child(current_cell)
# Then, we add the current block.
out.add_child(child)
current_cell = None
else:
out.add_child(child)
# Add the last current cell (if it had no output).
if current_cell:
out.add_child(current_cell)
return out
def ast():
# TODO: move this to conftest
ast = ASTNode('root')
# First block
block = ASTNode(name='Para', children=['hello '])
inline = ASTNode(name='Emph')
inline.add_child('world')
block.add_child(inline)
ast.add_child(block)
return ast
def ast():
# TODO: move this to conftest
ast = ASTNode('root')
# First block
block = ASTNode(name='Para',
children=['hello '])
inline = ASTNode(name='Emph')
inline.add_child('world')
block.add_child(inline)
ast.add_child(block)
return ast
