"""Read and write ``.vsk`` files.
A document is UTF-8 JSON, indented by two spaces, with a trailing newline,
so a ``.vsk`` diffs cleanly under version control. Writing is atomic: the
file is written beside its target and renamed over it, so an interrupted
save never leaves half a document.
"""
from __future__ import annotations
import json
import os
import stat
import tempfile
import threading
from pathlib import Path
from typing import Any
from pydantic import ValidationError
from vibesketch.document.model import Document, Edit, Provenance
__all__ = [
"ABSENT",
"DocumentError",
"SaveRefused",
"check_output",
"digest_of",
"dumps",
"load",
"loads",
"new_document",
"save",
"validation_errors",
"with_edit",
"write_output",
]
#: The file extension of a vibe-sketch document.
SUFFIX = ".vsk"
[docs]
class DocumentError(ValueError):
"""A document that is not valid JSON or not a valid vibe-sketch document.
``errors`` lists every problem as ``"location: message"``, the form the
CLI prints and the REST API returns.
"""
def __init__(self, errors: list[str]) -> None:
super().__init__("invalid vibe-sketch document:\n " + "\n ".join(errors))
self.errors = errors
[docs]
class SaveRefused(ValueError):
"""Writing would discard QVF result data or replace a changed target."""
def _reject_constant(value: str) -> None:
raise ValueError(f"{value} is not a JSON number")
[docs]
def validation_errors(error: ValidationError) -> list[str]:
"""Readable ``location: message`` lines for a pydantic validation error."""
lines = []
for item in error.errors(include_url=False):
location = ".".join(str(part) for part in item["loc"]) or "document"
lines.append(f"{location}: {item['msg']}")
return lines
def from_data(data: Any) -> Document:
"""Validate already-parsed JSON data as a document, within :mod:`vibesketch.limits`."""
from vibesketch.limits import LimitExceeded, check_depth
try:
check_depth(data, "the document")
except LimitExceeded as error:
raise DocumentError([f"document: {error}"]) from None
try:
return Document.model_validate(data)
except ValidationError as error:
raise DocumentError(validation_errors(error)) from error
[docs]
def loads(text: str) -> Document:
"""Parse and validate a document from JSON text, within :mod:`vibesketch.limits`."""
from vibesketch.limits import MAX_JSON_DEPTH, LimitExceeded, check_depth, check_size
try:
check_size(text, "the document")
data = json.loads(text, parse_constant=_reject_constant)
check_depth(data, "the document")
except RecursionError:
# Deeper than the parser follows, which varies with the Python version.
raise DocumentError([f"document: nests more than {MAX_JSON_DEPTH} levels deep"]) from None
except LimitExceeded as error:
raise DocumentError([f"document: {error}"]) from None
except ValueError as error:
raise DocumentError([f"document: not valid JSON ({error})"]) from error
return from_data(data)
[docs]
def load(path: str | os.PathLike[str]) -> Document:
"""Read and validate a document: a ``.qvf`` archive or a ``.vsk`` file."""
from vibesketch.limits import MAX_INPUT_BYTES, MAX_QVF_BYTES, LimitExceeded, read_file
source = Path(path)
qvf = source.suffix.lower() == ".qvf"
try:
content = read_file(source, source.name, limit=MAX_QVF_BYTES if qvf else MAX_INPUT_BYTES)
except LimitExceeded as error:
raise DocumentError([f"{source.name}: {error}"]) from None
if qvf:
from vibesketch.document.qvf import QvfError, from_qvf
try:
return from_qvf(content).document
except QvfError as error:
raise DocumentError([f"{source.name}: {error}"]) from None
return loads(content.decode("utf-8"))
[docs]
def dumps(document: Document) -> str:
"""The document as ``.vsk`` text."""
return json.dumps(document.to_json_data(), indent=2, ensure_ascii=False) + "\n"
_UMASK_LOCK = threading.Lock()
def _new_file_mode() -> int:
"""The mode a newly created file gets under the current umask."""
with _UMASK_LOCK: # reading the umask means setting it; keep threads out
umask = os.umask(0o022)
os.umask(umask)
return 0o666 & ~umask
#: What :func:`digest_of` says of a file that does not exist.
ABSENT = "absent"
_CHANGED = (
"The file changed outside vibe-sketch since it was opened or last saved, so saving "
"would discard that change. Reopen the file to see it, or keep this version with "
"Save As or a download under another name."
)
[docs]
def digest_of(path: str | os.PathLike[str]) -> str:
"""The SHA-256 of a file's bytes, or :data:`ABSENT` when there is no file."""
import hashlib
try:
digest = hashlib.sha256()
with Path(path).open("rb") as source:
for chunk in iter(lambda: source.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
except FileNotFoundError:
return ABSENT
[docs]
def save(document: Document, path: str | os.PathLike[str], *, expect: str | None = None) -> Path:
"""Write ``document`` to ``path`` atomically; returns the path written.
A ``.qvf`` path gets a QVF archive (raising
:class:`~vibesketch.document.qvf.QvfUnavailable` without qvf-writer), any
other a ``.vsk``. A file that exists keeps its permissions; a new one gets
the umask's. QVF results and archives with data outside the native sketch
raise :class:`SaveRefused`; save the sketch to another path instead.
With ``expect`` (a :func:`digest_of`), a file that no longer holds those
bytes is not replaced either.
"""
target = Path(path)
if target.suffix.lower() == ".qvf":
from vibesketch.document.qvf import to_qvf
content = to_qvf(document)
else:
content = dumps(document).encode("utf-8")
return write_output(target, content, expect=expect)
def _qvf_snapshot(target: Path) -> bytes | None:
from vibesketch.limits import MAX_QVF_BYTES, LimitExceeded, read_file
try:
return read_file(target, target.name, limit=MAX_QVF_BYTES)
except FileNotFoundError:
return None
except LimitExceeded as error:
raise SaveRefused(f"Cannot inspect the QVF destination: {error}. Use Save As.") from None
[docs]
def check_output(path: str | os.PathLike[str]) -> None:
"""Refuse a protected QVF destination before an operation with side effects.
:func:`write_output` checks again at the actual write; this preflight is
useful before submitting jobs whose run records must be saved in place.
"""
target = Path(path)
if ".qvf" in (target.suffix.lower(), target.resolve().suffix.lower()):
from vibesketch.document.qvf import overwrite_refusal
existing = _qvf_snapshot(target)
if existing is not None and (reason := overwrite_refusal(existing)):
raise SaveRefused(reason)
[docs]
def write_output(
path: str | os.PathLike[str], content: bytes, *, expect: str | None = None
) -> Path:
"""Atomically write exported bytes, protecting existing QVF result files.
This also guards explicit format overrides (such as SVG sent to a QVF
filename). Symlinks retain their link and write through to their target.
With ``expect``, the file must still hold the bytes that :func:`digest_of`
measured (or still be absent), before writing and again at the rename;
otherwise :class:`SaveRefused`, and the file is left as it is.
The new bytes reach the disk before they replace the old ones, and the
directory entry after, so a crash or a full disk leaves either the old file
or the new one, never a mixture. A crash at the wrong moment can leave a
hidden ``.NAME.*.tmp`` beside it, which is safe to delete.
"""
original = Path(path)
target = original.resolve()
if expect is not None and digest_of(target) != expect:
raise SaveRefused(_CHANGED)
guarded = ".qvf" in (original.suffix.lower(), target.suffix.lower())
previous = _qvf_snapshot(target) if guarded else None
if guarded and previous is not None:
from vibesketch.document.qvf import overwrite_refusal
if reason := overwrite_refusal(previous):
raise SaveRefused(reason)
mode = stat.S_IMODE(target.stat().st_mode) if target.exists() else _new_file_mode()
descriptor, temporary = tempfile.mkstemp(
prefix=f".{target.name}.", suffix=".tmp", dir=target.parent
)
try:
with os.fdopen(descriptor, "wb") as stream:
stream.write(content)
stream.flush()
os.fsync(stream.fileno())
os.chmod(temporary, mode)
if guarded and _qvf_snapshot(target) != previous:
raise SaveRefused("The QVF destination changed during saving. Use Save As.")
if expect is not None and digest_of(target) != expect:
raise SaveRefused(_CHANGED)
os.replace(temporary, target)
except BaseException:
Path(temporary).unlink(missing_ok=True)
raise
_sync_directory(target.parent)
return original
def _sync_directory(directory: Path) -> None:
"""Make a rename in ``directory`` durable; a no-op where directories cannot be synced."""
flags = getattr(os, "O_DIRECTORY", None)
if flags is None:
return
try:
descriptor = os.open(directory, os.O_RDONLY | flags)
except OSError:
return
try:
os.fsync(descriptor)
except OSError:
pass # some file systems refuse; the file itself is already synced
finally:
os.close(descriptor)
[docs]
def new_document(app: str | None = None) -> Document:
"""An empty document stamped with this application and the current time."""
if app is None:
from vibesketch import __version__
app = f"vibe-sketch {__version__}"
return Document(provenance=Provenance(app=app))
[docs]
def with_edit(document: Document, *, actor: str, client: str, summary: str) -> Document:
"""``document`` with an edit recorded for its next revision.
A run of alike edits (same actor, client and summary) folds into one
entry, so a drawing session does not leave hundreds. The revision itself
is the caller's to advance: the server's store does it on commit.
"""
from datetime import UTC, datetime
now = datetime.now(UTC).replace(microsecond=0)
edits = list(document.provenance.edits)
revision = document.revision + 1
if edits and (edits[-1].actor, edits[-1].client, edits[-1].summary) == (actor, client, summary):
edits[-1] = edits[-1].model_copy(update={"revision": revision, "at": now})
else:
fields = {"actor": actor, "client": client, "summary": summary}
edits.append(Edit.model_validate({"revision": revision, "at": now, **fields}))
provenance = document.provenance.model_copy(update={"edits": edits})
return document.model_copy(update={"provenance": provenance})