You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
221 lines
7.8 KiB
Python
221 lines
7.8 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 LICENSE, distributed as part of this software.
|
|
# -----------------------------------------------------------------------------
|
|
|
|
# -----------------------------------------------------------------------------
|
|
# Imports
|
|
# -----------------------------------------------------------------------------
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
from .nbbase import (
|
|
nbformat,
|
|
nbformat_minor,
|
|
new_code_cell,
|
|
new_heading_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):
|
|
"""An error raised for a pyreader error."""
|
|
|
|
|
|
class PyReader(NotebookReader):
|
|
"""A python notebook reader."""
|
|
|
|
def reads(self, s, **kwargs):
|
|
"""Convert a string to a notebook"""
|
|
return self.to_notebook(s, **kwargs)
|
|
|
|
def to_notebook(self, s, **kwargs):
|
|
"""Convert a string to a notebook"""
|
|
lines = s.splitlines()
|
|
cells = []
|
|
cell_lines: list[str] = []
|
|
kwargs = {}
|
|
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, **kwargs)
|
|
if cell is not None:
|
|
cells.append(cell)
|
|
state = "codecell"
|
|
cell_lines = []
|
|
kwargs = {}
|
|
elif line.startswith("# <htmlcell>"):
|
|
cell = self.new_cell(state, cell_lines, **kwargs)
|
|
if cell is not None:
|
|
cells.append(cell)
|
|
state = "htmlcell"
|
|
cell_lines = []
|
|
kwargs = {}
|
|
elif line.startswith("# <markdowncell>"):
|
|
cell = self.new_cell(state, cell_lines, **kwargs)
|
|
if cell is not None:
|
|
cells.append(cell)
|
|
state = "markdowncell"
|
|
cell_lines = []
|
|
kwargs = {}
|
|
# VERSIONHACK: plaintext -> raw
|
|
elif line.startswith(("# <rawcell>", "# <plaintextcell>")):
|
|
cell = self.new_cell(state, cell_lines, **kwargs)
|
|
if cell is not None:
|
|
cells.append(cell)
|
|
state = "rawcell"
|
|
cell_lines = []
|
|
kwargs = {}
|
|
elif line.startswith("# <headingcell"):
|
|
cell = self.new_cell(state, cell_lines, **kwargs)
|
|
if cell is not None:
|
|
cells.append(cell)
|
|
cell_lines = []
|
|
m = re.match(r"# <headingcell level=(?P<level>\d)>", line)
|
|
if m is not None:
|
|
state = "headingcell"
|
|
kwargs = {}
|
|
kwargs["level"] = int(m.group("level"))
|
|
else:
|
|
state = "codecell"
|
|
kwargs = {}
|
|
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)
|
|
return new_notebook(worksheets=[ws])
|
|
|
|
def new_cell(self, state, lines, **kwargs):
|
|
"""Create a new cell."""
|
|
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)
|
|
elif state == "rawcell":
|
|
text = self._remove_comments(lines)
|
|
if text:
|
|
return new_text_cell("raw", source=text)
|
|
elif state == "headingcell":
|
|
text = self._remove_comments(lines)
|
|
level = kwargs.get("level", 1)
|
|
if text:
|
|
return new_heading_cell(source=text, level=level)
|
|
|
|
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 # noqa: RET504
|
|
|
|
def split_lines_into_blocks(self, lines):
|
|
"""Split lines into code blocks."""
|
|
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):
|
|
"""A Python notebook writer."""
|
|
|
|
def writes(self, nb, **kwargs):
|
|
"""Convert a notebook to a string."""
|
|
lines = ["# -*- coding: utf-8 -*-"]
|
|
lines.extend(
|
|
[
|
|
"# <nbformat>%i.%i</nbformat>" % (nbformat, nbformat_minor),
|
|
"",
|
|
]
|
|
)
|
|
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("")
|
|
elif cell.cell_type == "raw":
|
|
input_ = cell.get("source")
|
|
if input_ is not None:
|
|
lines.extend(["# <rawcell>", ""])
|
|
lines.extend(["# " + line for line in input_.splitlines()])
|
|
lines.append("")
|
|
elif cell.cell_type == "heading":
|
|
input_ = cell.get("source")
|
|
level = cell.get("level", 1)
|
|
if input_ is not None:
|
|
lines.extend(["# <headingcell level=%s>" % level, ""])
|
|
lines.extend(["# " + line for line in input_.splitlines()])
|
|
lines.append("")
|
|
lines.append("")
|
|
return "\n".join(lines)
|
|
|
|
|
|
_reader = PyReader()
|
|
_writer = PyWriter()
|
|
|
|
reads = _reader.reads
|
|
read = _reader.read
|
|
to_notebook = _reader.to_notebook
|
|
write = _writer.write
|
|
writes = _writer.writes
|