2022-05-23 00:16:32 +04:00

152 lines
5.1 KiB
Python

"""Read and write notebooks as regular .py files.
Authors:
* Brian Granger
"""
# -----------------------------------------------------------------------------
# Copyright (C) 2008-2011 The IPython Development Team
#
# Distributed under the terms of the BSD License. The full license is in
# the file COPYING, distributed as part of this software.
# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------
# Imports
# -----------------------------------------------------------------------------
import re
from typing import List
from .nbbase import new_code_cell, new_notebook, new_text_cell, new_worksheet
from .rwbase import NotebookReader, NotebookWriter
# -----------------------------------------------------------------------------
# Code
# -----------------------------------------------------------------------------
_encoding_declaration_re = re.compile(r"^#.*coding[:=]\s*([-\w.]+)")
class PyReaderError(Exception):
pass
class PyReader(NotebookReader):
def reads(self, s, **kwargs):
return self.to_notebook(s, **kwargs)
def to_notebook(self, s, **kwargs):
lines = s.splitlines()
cells = []
cell_lines: List[str] = []
state = "codecell"
for line in lines:
if line.startswith("# <nbformat>") or _encoding_declaration_re.match(line):
pass
elif line.startswith("# <codecell>"):
cell = self.new_cell(state, cell_lines)
if cell is not None:
cells.append(cell)
state = "codecell"
cell_lines = []
elif line.startswith("# <htmlcell>"):
cell = self.new_cell(state, cell_lines)
if cell is not None:
cells.append(cell)
state = "htmlcell"
cell_lines = []
elif line.startswith("# <markdowncell>"):
cell = self.new_cell(state, cell_lines)
if cell is not None:
cells.append(cell)
state = "markdowncell"
cell_lines = []
else:
cell_lines.append(line)
if cell_lines and state == "codecell":
cell = self.new_cell(state, cell_lines)
if cell is not None:
cells.append(cell)
ws = new_worksheet(cells=cells)
nb = new_notebook(worksheets=[ws])
return nb
def new_cell(self, state, lines):
if state == "codecell":
input = "\n".join(lines)
input = input.strip("\n")
if input:
return new_code_cell(input=input)
elif state == "htmlcell":
text = self._remove_comments(lines)
if text:
return new_text_cell("html", source=text)
elif state == "markdowncell":
text = self._remove_comments(lines)
if text:
return new_text_cell("markdown", source=text)
def _remove_comments(self, lines):
new_lines = []
for line in lines:
if line.startswith("#"):
new_lines.append(line[2:])
else:
new_lines.append(line)
text = "\n".join(new_lines)
text = text.strip("\n")
return text
def split_lines_into_blocks(self, lines):
if len(lines) == 1:
yield lines[0]
raise StopIteration()
import ast
source = "\n".join(lines)
code = ast.parse(source)
starts = [x.lineno - 1 for x in code.body]
for i in range(len(starts) - 1):
yield "\n".join(lines[starts[i] : starts[i + 1]]).strip("\n")
yield "\n".join(lines[starts[-1] :]).strip("\n")
class PyWriter(NotebookWriter):
def writes(self, nb, **kwargs):
lines = ["# -*- coding: utf-8 -*-"]
lines.extend(["# <nbformat>2</nbformat>", ""])
for ws in nb.worksheets:
for cell in ws.cells:
if cell.cell_type == "code":
input = cell.get("input")
if input is not None:
lines.extend(["# <codecell>", ""])
lines.extend(input.splitlines())
lines.append("")
elif cell.cell_type == "html":
input = cell.get("source")
if input is not None:
lines.extend(["# <htmlcell>", ""])
lines.extend(["# " + line for line in input.splitlines()])
lines.append("")
elif cell.cell_type == "markdown":
input = cell.get("source")
if input is not None:
lines.extend(["# <markdowncell>", ""])
lines.extend(["# " + line for line in input.splitlines()])
lines.append("")
lines.append("")
return str("\n".join(lines))
_reader = PyReader()
_writer = PyWriter()
reads = _reader.reads
read = _reader.read
to_notebook = _reader.to_notebook
write = _writer.write
writes = _writer.writes