Repository Public Research preview
1"""Small, dependency-light I/O helpers (JSONL, atomic writes, Parquet)."""2from __future__ import annotations34import gzip5import json6import os7import tempfile8from pathlib import Path9from typing import Any, Iterable, Iterator101112def ensure_dir(p: Path | str) -> Path:13 p = Path(p)14 p.mkdir(parents=True, exist_ok=True)15 return p161718def atomic_write_text(path: Path | str, text: str, encoding: str = "utf-8") -> None:19 path = Path(path)20 ensure_dir(path.parent)21 fd, tmp = tempfile.mkstemp(dir=str(path.parent), prefix=".tmp-", suffix=path.suffix)22 try:23 with os.fdopen(fd, "w", encoding=encoding) as f:24 f.write(text)25 os.replace(tmp, path)26 finally:27 if os.path.exists(tmp):28 os.unlink(tmp)293031def atomic_write_bytes(path: Path | str, data: bytes) -> None:32 path = Path(path)33 ensure_dir(path.parent)34 fd, tmp = tempfile.mkstemp(dir=str(path.parent), prefix=".tmp-", suffix=path.suffix)35 try:36 with os.fdopen(fd, "wb") as f:37 f.write(data)38 os.replace(tmp, path)39 finally:40 if os.path.exists(tmp):41 os.unlink(tmp)424344def write_json(path: Path | str, obj: Any, indent: int = 2) -> None:45 atomic_write_text(path, json.dumps(obj, indent=indent, ensure_ascii=False, sort_keys=False, default=_default) + "\n")464748def read_json(path: Path | str) -> Any:49 with open(path, encoding="utf-8") as f:50 return json.load(f)515253def _default(o: Any) -> Any:54 if hasattr(o, "model_dump"):55 return o.model_dump()56 if isinstance(o, Path):57 return str(o)58 if isinstance(o, (set, frozenset)):59 return sorted(o)60 raise TypeError(f"not JSON serialisable: {type(o)}")616263def _open(path: Path | str, mode: str):64 path = Path(path)65 if path.suffix == ".gz":66 return gzip.open(path, mode + "t" if "b" not in mode else mode, encoding=None if "b" in mode else "utf-8")67 return open(path, mode, encoding="utf-8")686970def read_jsonl(path: Path | str) -> Iterator[dict]:71 path = Path(path)72 if not path.exists():73 return iter(())74 return _iter_jsonl(path)757677def _iter_jsonl(path: Path) -> Iterator[dict]:78 with _open(path, "r") as f:79 for line in f:80 line = line.strip()81 if line:82 yield json.loads(line)838485def write_jsonl(path: Path | str, rows: Iterable[Any], append: bool = False) -> int:86 path = Path(path)87 ensure_dir(path.parent)88 n = 089 with _open(path, "a" if append else "w") as f:90 for r in rows:91 f.write(json.dumps(r, ensure_ascii=False, default=_default) + "\n")92 n += 193 return n949596def append_jsonl(path: Path | str, row: Any) -> None:97 write_jsonl(path, [row], append=True)9899100def write_parquet(path: Path | str, rows: list[dict]) -> None:101 import pyarrow as pa102 import pyarrow.parquet as pq103104 ensure_dir(Path(path).parent)105 table = pa.Table.from_pylist(rows)106 pq.write_table(table, str(path))107108109def read_parquet(path: Path | str) -> list[dict]:110 import pyarrow.parquet as pq111112 return pq.read_table(str(path)).to_pylist()113