practize_1course/automation/clean_wolfram_notebook.py

660 lines
22 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""Удаляет из текстовых Wolfram Notebook (.nb) ячейки с Manipulate/CellPrint.
Удаляется вся вычислительная группа CellGroupData, если одна из её Input-ячеек
содержит Manipulate[...] или CellPrint[...]. Поэтому вместе с исходной ячейкой
исчезают связанные Output/Message/Print-ячейки.
Служебные блоки Notebook (CacheID, NotebookFileOutline, CellTagsIndex и весь
Internal cache information) не анализируются и не изменяются. Все нетронутые
части списка ячеек также сохраняются посимвольно.
Примеры:
python clean_wolfram_notebook.py notebook.nb
python clean_wolfram_notebook.py a.nb b.nb --output-dir cleaned
python clean_wolfram_notebook.py notebook.nb --in-place
python clean_wolfram_notebook.py notebook.nb --dry-run
"""
from __future__ import annotations
import argparse
import os
import re
import shutil
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Optional
RESULT_STYLES = {
"Output",
"Print",
"Message",
"Echo",
"GeneratedCell",
}
class NotebookFormatError(RuntimeError):
"""Файл не похож на корректный текстовый Wolfram Notebook."""
@dataclass
class Stats:
manipulate_groups: int = 0
cellprint_groups: int = 0
manipulate_cells: int = 0
cellprint_cells: int = 0
result_cells: int = 0
@property
def removed_groups(self) -> int:
return self.manipulate_groups + self.cellprint_groups
@property
def removed_input_cells(self) -> int:
return self.manipulate_cells + self.cellprint_cells
def add(self, other: "Stats") -> None:
self.manipulate_groups += other.manipulate_groups
self.cellprint_groups += other.cellprint_groups
self.manipulate_cells += other.manipulate_cells
self.cellprint_cells += other.cellprint_cells
self.result_cells += other.result_cells
@dataclass(frozen=True)
class ParsedCall:
open_index: int
close_index: int
args: list[str]
arg_ranges: list[tuple[int, int]]
@dataclass(frozen=True)
class ParsedList:
open_index: int
close_index: int
items: list[str]
item_ranges: list[tuple[int, int]]
@dataclass(frozen=True)
class InputTarget:
manipulate: bool = False
cellprint: bool = False
@property
def any(self) -> bool:
return self.manipulate or self.cellprint
def _scan_top_level_ranges(text: str, start: int, end: int) -> list[tuple[int, int]]:
"""Разбивает диапазон по запятым верхнего уровня.
Учитываются строки, экранирование, вложенные (), [] и {}, а также
вложенные комментарии Wolfram вида (* ... *).
"""
ranges: list[tuple[int, int]] = []
item_start = start
stack: list[str] = []
pairs = {")": "(", "]": "[", "}": "{"}
in_string = False
escaped = False
comment_depth = 0
i = start
while i < end:
ch = text[i]
nxt = text[i + 1] if i + 1 < end else ""
if comment_depth:
if ch == "(" and nxt == "*":
comment_depth += 1
i += 2
continue
if ch == "*" and nxt == ")":
comment_depth -= 1
i += 2
continue
i += 1
continue
if in_string:
if escaped:
escaped = False
elif ch == "\\":
escaped = True
elif ch == '"':
in_string = False
i += 1
continue
if ch == "(" and nxt == "*":
comment_depth = 1
i += 2
continue
if ch == '"':
in_string = True
i += 1
continue
if ch in "([{":
stack.append(ch)
i += 1
continue
if ch in ")]}" :
if not stack or stack[-1] != pairs[ch]:
raise NotebookFormatError(
f"Несогласованная скобка {ch!r} в позиции {i}"
)
stack.pop()
i += 1
continue
if ch == "," and not stack:
ranges.append((item_start, i))
item_start = i + 1
i += 1
if in_string:
raise NotebookFormatError("Незакрытая строка в Notebook")
if comment_depth:
raise NotebookFormatError("Незакрытый комментарий в Notebook")
if stack:
raise NotebookFormatError("Незакрытая скобка в Notebook")
ranges.append((item_start, end))
return ranges
def _find_matching(text: str, open_index: int) -> int:
opener = text[open_index]
closer = {"[": "]", "{": "}", "(": ")"}.get(opener)
if closer is None:
raise ValueError(f"Ожидалась открывающая скобка, получено {opener!r}")
stack: list[str] = [opener]
pairs = {")": "(", "]": "[", "}": "{"}
in_string = False
escaped = False
comment_depth = 0
i = open_index + 1
while i < len(text):
ch = text[i]
nxt = text[i + 1] if i + 1 < len(text) else ""
if comment_depth:
if ch == "(" and nxt == "*":
comment_depth += 1
i += 2
continue
if ch == "*" and nxt == ")":
comment_depth -= 1
i += 2
continue
i += 1
continue
if in_string:
if escaped:
escaped = False
elif ch == "\\":
escaped = True
elif ch == '"':
in_string = False
i += 1
continue
if ch == "(" and nxt == "*":
comment_depth = 1
i += 2
continue
if ch == '"':
in_string = True
i += 1
continue
if ch in "([{":
stack.append(ch)
i += 1
continue
if ch in ")]}" :
if not stack or stack[-1] != pairs[ch]:
raise NotebookFormatError(
f"Несогласованная скобка {ch!r} в позиции {i}"
)
stack.pop()
if not stack:
return i
i += 1
continue
i += 1
raise NotebookFormatError(
f"Не найдена закрывающая скобка {closer!r} для позиции {open_index}"
)
def _skip_space_and_comments(text: str, pos: int, end: Optional[int] = None) -> int:
limit = len(text) if end is None else end
while pos < limit:
if text[pos].isspace():
pos += 1
continue
if pos + 1 < limit and text[pos : pos + 2] == "(*":
depth = 1
pos += 2
while pos < limit and depth:
if pos + 1 < limit and text[pos : pos + 2] == "(*":
depth += 1
pos += 2
elif pos + 1 < limit and text[pos : pos + 2] == "*)":
depth -= 1
pos += 2
else:
pos += 1
if depth:
raise NotebookFormatError("Незакрытый комментарий")
continue
break
return pos
def _parse_call(expr: str, head: str) -> Optional[ParsedCall]:
pos = _skip_space_and_comments(expr, 0)
if not expr.startswith(head, pos):
return None
pos += len(head)
pos = _skip_space_and_comments(expr, pos)
if pos >= len(expr) or expr[pos] != "[":
return None
close = _find_matching(expr, pos)
tail = _skip_space_and_comments(expr, close + 1)
if tail != len(expr):
return None
ranges = _scan_top_level_ranges(expr, pos + 1, close)
return ParsedCall(
open_index=pos,
close_index=close,
args=[expr[a:b] for a, b in ranges],
arg_ranges=ranges,
)
def _parse_list(expr: str) -> Optional[ParsedList]:
pos = _skip_space_and_comments(expr, 0)
if pos >= len(expr) or expr[pos] != "{":
return None
close = _find_matching(expr, pos)
tail = _skip_space_and_comments(expr, close + 1)
if tail != len(expr):
return None
ranges = _scan_top_level_ranges(expr, pos + 1, close)
# Пустой список даёт один пустой диапазон — нормализуем его.
if len(ranges) == 1 and not expr[ranges[0][0] : ranges[0][1]].strip():
ranges = []
return ParsedList(
open_index=pos,
close_index=close,
items=[expr[a:b] for a, b in ranges],
item_ranges=ranges,
)
def _quoted_string(expr: str) -> Optional[str]:
s = expr.strip()
if len(s) < 2 or s[0] != '"' or s[-1] != '"':
return None
# Для стилей Notebook обычно нет escape-последовательностей.
return s[1:-1]
def _cell_style(cell_expr: str) -> Optional[str]:
cell = _parse_call(cell_expr.strip(), "Cell")
if cell is None or len(cell.args) < 2:
return None
return _quoted_string(cell.args[1])
def _input_target(cell_expr: str) -> InputTarget:
if _cell_style(cell_expr) != "Input":
return InputTarget()
# Стандартное коробочное представление: RowBox[{"Manipulate", "[", ...}]
manipulate = bool(
re.search(r'"Manipulate"\s*,\s*"\\?\["', cell_expr)
or re.search(r'(?<![A-Za-z0-9_$`])Manipulate\s*\[', cell_expr)
)
cellprint = bool(
re.search(r'"CellPrint"\s*,\s*"\\?\["', cell_expr)
or re.search(r'(?<![A-Za-z0-9_$`])CellPrint\s*\[', cell_expr)
)
return InputTarget(manipulate=manipulate, cellprint=cellprint)
def _join_modified_items(original_content: str, items: Iterable[str]) -> str:
"""Соединяет оставшиеся элементы, не форматируя их заново.
Каждый элемент передаётся как точный срез исходного файла вместе с его
пробелами и переводами строк. Меняются только удалённые элементы и
разделяющие их запятые.
"""
kept = [item for item in items if item.strip()]
if not kept:
# Пустое содержимое списка допустимо. Намеренно не переносим сюда
# пробелы удалённых ячеек, чтобы не оставлять огромные пустые области.
return ""
return ",".join(kept)
def _clean_cell_group(cell_expr: str) -> tuple[Optional[str], Stats]:
stats = Stats()
cell = _parse_call(cell_expr.strip(), "Cell")
if cell is None or not cell.args:
return cell_expr, stats
group_arg = cell.args[0]
group = _parse_call(group_arg.strip(), "CellGroupData")
if group is None or not group.args:
return cell_expr, stats
children_expr = group.args[0]
children_stripped = children_expr.strip()
children = _parse_list(children_stripped)
if children is None:
return cell_expr, stats
# Если непосредственно в этой вычислительной группе есть целевой Input,
# удаляем всю группу: так гарантированно уходят все её Output/Message/Print.
direct_target = InputTarget()
for child in children.items:
target = _input_target(child.strip())
direct_target = InputTarget(
manipulate=direct_target.manipulate or target.manipulate,
cellprint=direct_target.cellprint or target.cellprint,
)
if direct_target.any:
if direct_target.manipulate:
stats.manipulate_groups += 1
elif direct_target.cellprint:
stats.cellprint_groups += 1
# В редком случае обе конструкции находятся в одной группе, считаем
# группу один раз, с приоритетом Manipulate, но ячейки ниже не нужны.
return None, stats
cleaned_items, child_stats, changed = _clean_cell_list(
children_stripped[children.open_index + 1 : children.close_index]
)
stats.add(child_stats)
if not changed:
return cell_expr, stats
new_children_stripped = (
children_stripped[: children.open_index + 1]
+ cleaned_items
+ children_stripped[children.close_index :]
)
children_leading = children_expr[: len(children_expr) - len(children_expr.lstrip())]
children_trailing = children_expr[len(children_expr.rstrip()) :]
new_children_expr = children_leading + new_children_stripped + children_trailing
# Заменяем первый аргумент CellGroupData внутри исходного выражения.
group_stripped = group_arg.strip()
group_parsed = _parse_call(group_stripped, "CellGroupData")
assert group_parsed is not None
arg0_start, arg0_end = group_parsed.arg_ranges[0]
new_group_stripped = (
group_stripped[:arg0_start]
+ new_children_expr
+ group_stripped[arg0_end:]
)
group_leading = group_arg[: len(group_arg) - len(group_arg.lstrip())]
group_trailing = group_arg[len(group_arg.rstrip()) :]
new_group_arg = group_leading + new_group_stripped + group_trailing
cell_stripped = cell_expr.strip()
cell_parsed = _parse_call(cell_stripped, "Cell")
assert cell_parsed is not None
cell_arg0_start, cell_arg0_end = cell_parsed.arg_ranges[0]
new_cell_stripped = (
cell_stripped[:cell_arg0_start]
+ new_group_arg
+ cell_stripped[cell_arg0_end:]
)
leading = cell_expr[: len(cell_expr) - len(cell_expr.lstrip())]
trailing = cell_expr[len(cell_expr.rstrip()) :]
return leading + new_cell_stripped + trailing, stats
def _clean_cell_list(content: str) -> tuple[str, Stats, bool]:
ranges = _scan_top_level_ranges(content, 0, len(content))
if len(ranges) == 1 and not content[ranges[0][0] : ranges[0][1]].strip():
return content, Stats(), False
stats = Stats()
kept: list[str] = []
changed = False
remove_following_results = False
for start, end in ranges:
raw_item = content[start:end]
item = raw_item.strip()
if not item:
continue
style = _cell_style(item)
if remove_following_results:
if style in RESULT_STYLES:
stats.result_cells += 1
changed = True
continue
remove_following_results = False
target = _input_target(item)
if target.any:
if target.manipulate:
stats.manipulate_cells += 1
elif target.cellprint:
stats.cellprint_cells += 1
changed = True
remove_following_results = True
continue
cleaned_group, group_stats = _clean_cell_group(raw_item)
stats.add(group_stats)
if cleaned_group is None:
changed = True
continue
if cleaned_group != raw_item:
changed = True
kept.append(cleaned_group)
if not changed:
return content, stats, False
return _join_modified_items(content, kept), stats, True
def _find_notebook_cell_list(text: str) -> tuple[int, int]:
marker = "(* Beginning of Notebook Content *)"
search_from = text.find(marker)
if search_from < 0:
search_from = 0
match = re.search(r"\bNotebook\s*\[", text[search_from:])
if not match:
raise NotebookFormatError("Не найдено выражение Notebook[...]")
notebook_start = search_from + match.start()
bracket = text.find("[", notebook_start)
pos = _skip_space_and_comments(text, bracket + 1)
if pos >= len(text) or text[pos] != "{":
raise NotebookFormatError("Первый аргумент Notebook не является списком ячеек")
close = _find_matching(text, pos)
return pos, close
def clean_notebook_text(text: str) -> tuple[str, Stats, bool]:
"""Очищает только первый аргумент Notebook[{...}], не трогая кэш.
Префикс до списка ячеек и суффикс после него вставляются в результат без
каких-либо преобразований. Это гарантирует сохранность CacheID и всего
блока Internal cache information.
"""
list_open, list_close = _find_notebook_cell_list(text)
prefix = text[: list_open + 1]
original_content = text[list_open + 1 : list_close]
suffix = text[list_close:]
cleaned_content, stats, changed = _clean_cell_list(original_content)
cleaned = prefix + cleaned_content + suffix if changed else text
# Проверяем лишь синтаксическую доступность Notebook и списка ячеек.
# Служебные блоки после Notebook[...] намеренно не читаются и не меняются.
_find_notebook_cell_list(cleaned)
return cleaned, stats, changed
def _read_text(path: Path) -> tuple[str, str, bool]:
data = path.read_bytes()
has_bom = data.startswith(b"\xef\xbb\xbf")
if has_bom:
data = data[3:]
try:
text = data.decode("utf-8")
except UnicodeDecodeError as exc:
raise NotebookFormatError(
f"{path}: ожидается текстовый .nb в UTF-8/ASCII: {exc}"
) from exc
newline = "\r\n" if b"\r\n" in data else "\n"
text = text.replace("\r\n", "\n").replace("\r", "\n")
return text, newline, has_bom
def _encode_text(text: str, newline: str, has_bom: bool) -> bytes:
if newline != "\n":
text = text.replace("\n", newline)
data = text.encode("utf-8")
return (b"\xef\xbb\xbf" + data) if has_bom else data
def _atomic_write(path: Path, data: bytes, mode: int = 0o644) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
fd, temp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
temp_path = Path(temp_name)
try:
with os.fdopen(fd, "wb") as fh:
fh.write(data)
fh.flush()
os.fsync(fh.fileno())
os.chmod(temp_path, mode)
os.replace(temp_path, path)
except Exception:
temp_path.unlink(missing_ok=True)
raise
def _default_output_path(input_path: Path) -> Path:
return input_path.with_name(f"{input_path.stem}_cleaned{input_path.suffix}")
def _resolve_outputs(args: argparse.Namespace, inputs: list[Path]) -> list[Path]:
if args.output:
if len(inputs) != 1:
raise SystemExit("Параметр --output можно использовать только с одним входным файлом")
return [Path(args.output)]
if args.output_dir:
out_dir = Path(args.output_dir)
return [out_dir / f"{p.stem}_cleaned{p.suffix}" for p in inputs]
if args.in_place:
return inputs
return [_default_output_path(p) for p in inputs]
def _print_stats(path: Path, stats: Stats, changed: bool, dry_run: bool) -> None:
action = "Проверен" if dry_run else "Готов"
status = "изменения найдены" if changed else "целевых ячеек нет"
print(
f"{action}: {path}{status}; "
f"групп Manipulate: {stats.manipulate_groups}, "
f"групп CellPrint: {stats.cellprint_groups}, "
f"отдельных Input: {stats.removed_input_cells}, "
f"отдельных результатов: {stats.result_cells}"
)
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description=(
"Удаляет из Wolfram Notebook группы с Manipulate/CellPrint "
"вместе с их результатами."
)
)
parser.add_argument("inputs", nargs="+", help="входные .nb-файлы")
out = parser.add_mutually_exclusive_group()
out.add_argument("-o", "--output", help="выходной файл (только для одного входа)")
out.add_argument("--output-dir", help="каталог для очищенных файлов")
out.add_argument(
"--in-place",
action="store_true",
help="перезаписать входные файлы; рядом создаётся резервная копия .bak",
)
parser.add_argument(
"--no-backup",
action="store_true",
help="не создавать .bak при --in-place",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="только показать, что будет удалено, не записывая файлы",
)
return parser
def main(argv: Optional[list[str]] = None) -> int:
parser = build_parser()
args = parser.parse_args(argv)
inputs = [Path(p) for p in args.inputs]
for path in inputs:
if not path.is_file():
parser.error(f"файл не найден: {path}")
if path.suffix.lower() != ".nb":
parser.error(f"ожидался файл .nb: {path}")
outputs = _resolve_outputs(args, inputs)
exit_code = 0
for input_path, output_path in zip(inputs, outputs):
try:
text, newline, has_bom = _read_text(input_path)
cleaned, stats, changed = clean_notebook_text(text)
if args.dry_run:
_print_stats(input_path, stats, changed, dry_run=True)
continue
if args.in_place and changed and not args.no_backup:
backup = input_path.with_suffix(input_path.suffix + ".bak")
shutil.copy2(input_path, backup)
input_mode = input_path.stat().st_mode & 0o777
_atomic_write(
output_path,
_encode_text(cleaned, newline, has_bom),
mode=input_mode or 0o644,
)
_print_stats(output_path, stats, changed, dry_run=False)
except (OSError, NotebookFormatError) as exc:
print(f"Ошибка: {input_path}: {exc}", file=sys.stderr)
exit_code = 1
return exit_code
if __name__ == "__main__":
raise SystemExit(main())