541 lines
19 KiB
Python
Executable file
541 lines
19 KiB
Python
Executable file
#!/usr/bin/env python3
|
||
"""
|
||
insert_d_into_nb.py
|
||
|
||
Рекурсивно обходит Wolfram Notebook (*.nb), пропускает файлы с
|
||
"фитирование" в имени и меняет значение d в первой Input-ячейке
|
||
на значение из файла коэффициентов.
|
||
|
||
Ключевая мера безопасности:
|
||
- notebook читается и пишется как bytes, без перекодирования и нормализации EOL;
|
||
- Internal cache information не удаляется и не редактируется;
|
||
- размер изменяемой первой Input-ячейки и размер всего .nb сохраняются
|
||
побайтно неизменными: если новое число длиннее/короче старого,
|
||
разница компенсируется незначащими пробелами в этой же Input-ячейке;
|
||
- перед любой записью создаётся backup в уникальной папке /tmp;
|
||
- запись атомарная, после неё выполняется повторная проверка.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import os
|
||
import re
|
||
import shutil
|
||
import stat
|
||
import sys
|
||
import tempfile
|
||
from dataclasses import dataclass
|
||
from decimal import Decimal, InvalidOperation
|
||
from pathlib import Path
|
||
|
||
|
||
SECTION_RE = re.compile(r"^\s*(\d+-\d+_\d+-\d+)\s*$")
|
||
D_LINE_RE = re.compile(
|
||
r"^\s*d\s*=\s*([+-]?(?:\d+(?:\.\d*)?|\.\d+))\s*$"
|
||
)
|
||
|
||
# Примеры:
|
||
# h_154-86_Pu240_8-9_gid_Qw1.nb
|
||
# h_152-88_h_Pu240_6-7_pr_Qw1.nb
|
||
FILENAME_KEY_RE = re.compile(
|
||
r"(?P<ab>\d+-\d+)(?:_h)?_Pu240_(?P<ij>\d+-\d+)",
|
||
re.IGNORECASE,
|
||
)
|
||
|
||
BEGIN_MARKER = b"(* Beginning of Notebook Content *)"
|
||
END_MARKER = b"(* End of Notebook Content *)"
|
||
|
||
# Ищем именно Cell[BoxData[...], "Input"] и не разрешаем матчу
|
||
# перескочить через начало следующей Cell[...].
|
||
FIRST_INPUT_RE = re.compile(
|
||
rb'Cell\[\s*BoxData\[\s*'
|
||
rb'(?P<box>(?:(?!\bCell\[).)*?)'
|
||
rb'\]\s*,\s*"Input"',
|
||
re.DOTALL,
|
||
)
|
||
|
||
# Для первой переменной принимается консервативная и ожидаемая по ТЗ форма:
|
||
# RowBox[{"d", "=", "2.88"}]
|
||
DIRECT_ASSIGNMENT_RE = re.compile(
|
||
rb'^\s*RowBox\[\{\s*'
|
||
rb'"(?P<var>[A-Za-z$][A-Za-z0-9$]*)"\s*,\s*'
|
||
rb'"="\s*,\s*'
|
||
rb'"(?P<value>[+-]?(?:\d+(?:\.\d*)?|\.\d+))"\s*'
|
||
rb'\}\]\s*$',
|
||
re.DOTALL,
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class FirstAssignment:
|
||
variable: str
|
||
value_text: str
|
||
value_start: int
|
||
value_end: int
|
||
first_input_end: int
|
||
next_cell_start: int
|
||
|
||
|
||
class NotebookFormatError(RuntimeError):
|
||
pass
|
||
|
||
|
||
def parse_args() -> argparse.Namespace:
|
||
parser = argparse.ArgumentParser(
|
||
description=(
|
||
"Заменяет d в первой Input-ячейке Wolfram Notebook "
|
||
"по значениям из файла коэффициентов."
|
||
)
|
||
)
|
||
parser.add_argument(
|
||
"--coeff_file",
|
||
required=True,
|
||
type=Path,
|
||
help="Путь к txt-файлу коэффициентов.",
|
||
)
|
||
parser.add_argument(
|
||
"--wolfram_notebooks",
|
||
"--wolfram_noteooks", # алиас для опечатки, встречающейся в ТЗ
|
||
dest="wolfram_notebooks",
|
||
required=True,
|
||
type=Path,
|
||
help="Корневая директория, где рекурсивно искать *.nb.",
|
||
)
|
||
return parser.parse_args()
|
||
|
||
|
||
def parse_coefficients(path: Path) -> dict[str, str]:
|
||
try:
|
||
text = path.read_text(encoding="utf-8-sig")
|
||
except OSError as exc:
|
||
raise RuntimeError(f"Не удалось прочитать файл коэффициентов: {exc}") from exc
|
||
|
||
result: dict[str, str] = {}
|
||
current_section: str | None = None
|
||
seen_sections: set[str] = set()
|
||
|
||
for lineno, raw_line in enumerate(text.splitlines(), start=1):
|
||
section_match = SECTION_RE.match(raw_line)
|
||
if section_match:
|
||
current_section = section_match.group(1)
|
||
if current_section in seen_sections:
|
||
raise RuntimeError(
|
||
f"{path}:{lineno}: повторная секция {current_section!r}"
|
||
)
|
||
seen_sections.add(current_section)
|
||
continue
|
||
|
||
d_match = D_LINE_RE.match(raw_line)
|
||
if d_match and current_section is not None:
|
||
if current_section in result:
|
||
raise RuntimeError(
|
||
f"{path}:{lineno}: повторное d в секции {current_section!r}"
|
||
)
|
||
value = d_match.group(1)
|
||
|
||
# Дополнительная валидация числа.
|
||
try:
|
||
Decimal(value)
|
||
except InvalidOperation as exc:
|
||
raise RuntimeError(
|
||
f"{path}:{lineno}: некорректное значение d={value!r}"
|
||
) from exc
|
||
|
||
result[current_section] = value
|
||
|
||
if not result:
|
||
raise RuntimeError(
|
||
f"В {path} не найдено ни одной секции с корректной строкой d = ..."
|
||
)
|
||
|
||
missing_d = sorted(seen_sections - result.keys())
|
||
if missing_d:
|
||
raise RuntimeError(
|
||
"В следующих секциях отсутствует корректная строка d: "
|
||
+ ", ".join(missing_d)
|
||
)
|
||
|
||
return result
|
||
|
||
|
||
def coefficient_key_from_filename(name: str) -> str | None:
|
||
match = FILENAME_KEY_RE.search(name)
|
||
if not match:
|
||
return None
|
||
return f"{match.group('ab')}_{match.group('ij')}"
|
||
|
||
|
||
def find_first_assignment(data: bytes) -> FirstAssignment:
|
||
begin = data.find(BEGIN_MARKER)
|
||
end = data.find(END_MARKER)
|
||
|
||
if begin < 0 or end < 0 or end <= begin:
|
||
raise NotebookFormatError(
|
||
"не найдены корректные маркеры начала/конца Notebook Content"
|
||
)
|
||
|
||
content_start = begin + len(BEGIN_MARKER)
|
||
content = data[content_start:end]
|
||
|
||
input_match = FIRST_INPUT_RE.search(content)
|
||
if not input_match:
|
||
raise NotebookFormatError("не найдена первая Input-ячейка")
|
||
|
||
box = input_match.group("box")
|
||
assignment = DIRECT_ASSIGNMENT_RE.fullmatch(box)
|
||
if not assignment:
|
||
raise NotebookFormatError(
|
||
"первая Input-ячейка не является простым присваиванием вида "
|
||
'RowBox[{"d", "=", "число"}]'
|
||
)
|
||
|
||
var = assignment.group("var").decode("ascii")
|
||
value = assignment.group("value").decode("ascii")
|
||
|
||
box_abs_start = content_start + input_match.start("box")
|
||
value_start = box_abs_start + assignment.start("value")
|
||
value_end = box_abs_start + assignment.end("value")
|
||
first_input_end = content_start + input_match.end()
|
||
|
||
# Компенсацию длины делаем до следующей Cell[, то есть внутри
|
||
# опций первой Input-ячейки и до начала следующей ячейки.
|
||
next_cell_start = data.find(b"Cell[", first_input_end, end)
|
||
if next_cell_start < 0:
|
||
raise NotebookFormatError(
|
||
"после первой Input-ячейки не найдена следующая Cell"
|
||
)
|
||
|
||
return FirstAssignment(
|
||
variable=var,
|
||
value_text=value,
|
||
value_start=value_start,
|
||
value_end=value_end,
|
||
first_input_end=first_input_end,
|
||
next_cell_start=next_cell_start,
|
||
)
|
||
|
||
|
||
def decimal_equal(left: str, right: str) -> bool:
|
||
try:
|
||
return Decimal(left) == Decimal(right)
|
||
except InvalidOperation:
|
||
return left == right
|
||
|
||
|
||
def _remove_indentation_spaces(data: bytes, count: int) -> bytes:
|
||
"""
|
||
Удаляет ровно count пробелов/табов только из отступов в начале строк.
|
||
|
||
Переводы строк не меняются, содержимое строковых литералов не трогается.
|
||
Для типичного CellChangeTimes в первой ячейке таких пробелов значительно
|
||
больше, чем может понадобиться для компенсации длины d.
|
||
"""
|
||
if count <= 0:
|
||
return data
|
||
|
||
out = bytearray()
|
||
i = 0
|
||
remaining = count
|
||
at_line_start = False
|
||
|
||
while i < len(data):
|
||
b = data[i]
|
||
|
||
if at_line_start and b in (0x20, 0x09) and remaining > 0:
|
||
# Удаляем отступ по одному байту.
|
||
remaining -= 1
|
||
i += 1
|
||
continue
|
||
|
||
out.append(b)
|
||
|
||
if b == 0x0A: # \n; CR из CRLF уже был скопирован
|
||
at_line_start = True
|
||
elif b not in (0x20, 0x09, 0x0D):
|
||
at_line_start = False
|
||
|
||
i += 1
|
||
|
||
if remaining:
|
||
raise NotebookFormatError(
|
||
"не хватило безопасных отступов для сохранения размера notebook "
|
||
f"(нужно было убрать ещё {remaining} байт)"
|
||
)
|
||
|
||
return bytes(out)
|
||
|
||
|
||
def patch_d_preserving_size(
|
||
original: bytes,
|
||
assignment: FirstAssignment,
|
||
new_value: str,
|
||
) -> bytes:
|
||
new_value_bytes = new_value.encode("ascii")
|
||
old_value_bytes = original[assignment.value_start:assignment.value_end]
|
||
|
||
if old_value_bytes.decode("ascii") != assignment.value_text:
|
||
raise NotebookFormatError("внутренняя проверка позиции старого d не прошла")
|
||
|
||
before = original[:assignment.value_start]
|
||
after_value = original[assignment.value_end:assignment.next_cell_start]
|
||
tail = original[assignment.next_cell_start:]
|
||
|
||
delta = len(new_value_bytes) - len(old_value_bytes)
|
||
|
||
if delta > 0:
|
||
# Новое число длиннее. Убираем столько же незначащих пробелов
|
||
# из отступов ПОСЛЕ значения d, но до следующей Cell.
|
||
compensated_after = _remove_indentation_spaces(after_value, delta)
|
||
elif delta < 0:
|
||
# Новое число короче. Добавляем пробелы сразу ПОСЛЕ закрывающей
|
||
# кавычки числового RowBox-токена. Это Wolfram-синтаксически
|
||
# незначащий whitespace.
|
||
need = -delta
|
||
if not after_value.startswith(b'"'):
|
||
raise NotebookFormatError(
|
||
"неожиданная структура после числового токена d"
|
||
)
|
||
compensated_after = (
|
||
after_value[:1] + (b" " * need) + after_value[1:]
|
||
)
|
||
else:
|
||
compensated_after = after_value
|
||
|
||
patched = before + new_value_bytes + compensated_after + tail
|
||
|
||
if len(patched) != len(original):
|
||
raise NotebookFormatError(
|
||
"после компенсации изменился размер notebook — запись отменена"
|
||
)
|
||
|
||
# Всё после начала следующей Cell должно остаться байт-в-байт неизменным.
|
||
if patched[assignment.next_cell_start:] != original[assignment.next_cell_start:]:
|
||
raise NotebookFormatError(
|
||
"обнаружено изменение за пределами первой Input-ячейки"
|
||
)
|
||
|
||
# Cache-блок должен остаться полностью неизменным.
|
||
old_cache_pos = original.find(b"(* Internal cache information *)", end_marker_pos(original))
|
||
new_cache_pos = patched.find(b"(* Internal cache information *)", end_marker_pos(patched))
|
||
if old_cache_pos < 0 or new_cache_pos < 0:
|
||
raise NotebookFormatError("не найден завершающий Internal cache information")
|
||
if old_cache_pos != new_cache_pos:
|
||
raise NotebookFormatError("сместилась позиция завершающего cache-блока")
|
||
if patched[old_cache_pos:] != original[old_cache_pos:]:
|
||
raise NotebookFormatError("cache-блок был изменён — запись отменена")
|
||
|
||
# Повторно разбираем результат и убеждаемся, что d стал целевым.
|
||
check = find_first_assignment(patched)
|
||
if check.variable != "d" or check.value_text != new_value:
|
||
raise NotebookFormatError(
|
||
"контрольная проверка после замены d не прошла"
|
||
)
|
||
|
||
return patched
|
||
|
||
|
||
def end_marker_pos(data: bytes) -> int:
|
||
pos = data.find(END_MARKER)
|
||
return pos if pos >= 0 else 0
|
||
|
||
|
||
def backup_file(source: Path, root: Path, backup_root: Path) -> Path:
|
||
relative = source.relative_to(root)
|
||
destination = backup_root / relative
|
||
destination.parent.mkdir(parents=True, exist_ok=True)
|
||
shutil.copy2(source, destination)
|
||
return destination
|
||
|
||
|
||
def atomic_write(path: Path, data: bytes) -> None:
|
||
original_stat = path.stat()
|
||
temp_name: str | None = None
|
||
|
||
try:
|
||
with tempfile.NamedTemporaryFile(
|
||
mode="wb",
|
||
prefix=f".{path.name}.",
|
||
suffix=".tmp",
|
||
dir=path.parent,
|
||
delete=False,
|
||
) as tmp:
|
||
temp_name = tmp.name
|
||
tmp.write(data)
|
||
tmp.flush()
|
||
os.fsync(tmp.fileno())
|
||
|
||
os.chmod(temp_name, stat.S_IMODE(original_stat.st_mode))
|
||
os.replace(temp_name, path)
|
||
temp_name = None
|
||
|
||
# На Linux fsync директории закрепляет rename в файловой системе.
|
||
try:
|
||
dir_fd = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY)
|
||
except OSError:
|
||
dir_fd = None
|
||
if dir_fd is not None:
|
||
try:
|
||
os.fsync(dir_fd)
|
||
finally:
|
||
os.close(dir_fd)
|
||
|
||
finally:
|
||
if temp_name is not None:
|
||
try:
|
||
os.unlink(temp_name)
|
||
except FileNotFoundError:
|
||
pass
|
||
|
||
|
||
def process_notebook(
|
||
path: Path,
|
||
root: Path,
|
||
coefficients: dict[str, str],
|
||
backup_root: Path,
|
||
) -> str:
|
||
if "фитирование" in path.name.casefold():
|
||
print(f"[SKIP fitting] {path}")
|
||
return "skipped_fitting"
|
||
|
||
if path.is_symlink():
|
||
print(f"[SKIP symlink] {path}")
|
||
return "skipped_other"
|
||
|
||
key = coefficient_key_from_filename(path.name)
|
||
if key is None:
|
||
print(f"[SKIP no-key] {path}: не удалось определить секцию коэффициентов")
|
||
return "skipped_other"
|
||
|
||
target = coefficients.get(key)
|
||
if target is None:
|
||
print(f"[SKIP no-coeff] {path}: нет секции {key!r} в coeff_file")
|
||
return "skipped_other"
|
||
|
||
try:
|
||
original = path.read_bytes()
|
||
except OSError as exc:
|
||
print(f"[ERROR read] {path}: {exc}")
|
||
return "error"
|
||
|
||
try:
|
||
first = find_first_assignment(original)
|
||
except NotebookFormatError as exc:
|
||
print(f"[SKIP first-d] {path}: {exc}")
|
||
return "skipped_first_d"
|
||
|
||
if first.variable != "d":
|
||
print(
|
||
f"[SKIP first-d] {path}: d не первая переменная "
|
||
f"(первая переменная: {first.variable})"
|
||
)
|
||
return "skipped_first_d"
|
||
|
||
if decimal_equal(first.value_text, target):
|
||
print(f"[OK already] {path}: d = {first.value_text}")
|
||
return "already"
|
||
|
||
try:
|
||
patched = patch_d_preserving_size(original, first, target)
|
||
except NotebookFormatError as exc:
|
||
print(f"[ERROR patch] {path}: {exc}")
|
||
return "error"
|
||
|
||
try:
|
||
backup_path = backup_file(path, root, backup_root)
|
||
except OSError as exc:
|
||
print(f"[ERROR backup] {path}: {exc}")
|
||
return "error"
|
||
|
||
try:
|
||
atomic_write(path, patched)
|
||
|
||
# Проверяем фактически записанные байты, а не только буфер.
|
||
written = path.read_bytes()
|
||
if written != patched:
|
||
raise OSError("байты после записи не совпадают с проверенным буфером")
|
||
|
||
verified = find_first_assignment(written)
|
||
if verified.variable != "d" or verified.value_text != target:
|
||
raise OSError("d после записи не прошёл контрольную проверку")
|
||
|
||
except Exception as exc:
|
||
# Если что-то пошло не так после backup, откатываем исходник.
|
||
try:
|
||
shutil.copy2(backup_path, path)
|
||
restore_note = "исходный файл восстановлен из backup"
|
||
except OSError as restore_exc:
|
||
restore_note = f"ОШИБКА ВОССТАНОВЛЕНИЯ: {restore_exc}"
|
||
|
||
print(f"[ERROR write] {path}: {exc}; {restore_note}")
|
||
return "error"
|
||
|
||
print(
|
||
f"[CHANGED] {path}: d {first.value_text} -> {target} "
|
||
f"(размер файла сохранён: {len(original)} байт)"
|
||
)
|
||
return "changed"
|
||
|
||
|
||
def main() -> int:
|
||
args = parse_args()
|
||
|
||
coeff_file = args.coeff_file.expanduser().resolve()
|
||
root = args.wolfram_notebooks.expanduser().resolve()
|
||
|
||
if not coeff_file.is_file():
|
||
print(f"ERROR: coeff_file не является файлом: {coeff_file}", file=sys.stderr)
|
||
return 2
|
||
|
||
if not root.is_dir():
|
||
print(
|
||
f"ERROR: wolfram_notebooks не является директорией: {root}",
|
||
file=sys.stderr,
|
||
)
|
||
return 2
|
||
|
||
try:
|
||
coefficients = parse_coefficients(coeff_file)
|
||
except RuntimeError as exc:
|
||
print(f"ERROR: {exc}", file=sys.stderr)
|
||
return 2
|
||
|
||
# Папка специально НЕ удаляется после завершения.
|
||
backup_root = Path(
|
||
tempfile.mkdtemp(prefix="insert_d_into_nb_backup_", dir="/tmp")
|
||
)
|
||
|
||
counters = {
|
||
"changed": 0,
|
||
"already": 0,
|
||
"skipped_fitting": 0,
|
||
"skipped_first_d": 0,
|
||
"skipped_other": 0,
|
||
"error": 0,
|
||
}
|
||
|
||
notebooks = sorted(
|
||
(p for p in root.rglob("*.nb") if p.is_file() or p.is_symlink()),
|
||
key=lambda p: str(p),
|
||
)
|
||
|
||
print(f"Найдено notebook-файлов: {len(notebooks)}")
|
||
|
||
for path in notebooks:
|
||
result = process_notebook(path, root, coefficients, backup_root)
|
||
counters[result] += 1
|
||
|
||
print("\nИтог:")
|
||
print(f" изменено: {counters['changed']}")
|
||
print(f" уже имели правильное d: {counters['already']}")
|
||
print(f" пропущено (фитирование): {counters['skipped_fitting']}")
|
||
print(f" пропущено (d не первая): {counters['skipped_first_d']}")
|
||
print(f" пропущено (прочее): {counters['skipped_other']}")
|
||
print(f" ошибок: {counters['error']}")
|
||
print(f"Папка backup: {backup_root}")
|
||
|
||
return 1 if counters["error"] else 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|