null: packfreeze — deterministic blake2b package manifest generator with verify/check
This commit is contained in:
null committed
1 parent
5f789397b8
commit
5ca0b6c895
5 files changed
+1969
No files matched your search
@@ -0,0 +1,457 @@
|
||||
#!/usr/bin/env python3
|
||||
"""chunktext.py -- token-budgeted chunking with an unambiguous overlap carrier.
|
||||
|
||||
Stdlib only. No network. Deterministic. Single file.
|
||||
|
||||
Why this file exists
|
||||
--------------------
|
||||
``chunker.py`` (wizard-outputs/null) claims "lossless round-trip" and claims it
|
||||
escapes its own sentinel. Both claims are false as measured by
|
||||
``test_chunktext.py``:
|
||||
|
||||
* ``reconstruct`` cannot distinguish overlap from content, because the
|
||||
sentinel-delimited tail is never placed at a guaranteed boundary -- when a
|
||||
chunk ends in a word that does not start a sentence, the tail ends
|
||||
mid-word, so the join drops a letter. Measured: 74/120 round-trip cases
|
||||
differed from the input.
|
||||
* ``chunk_text`` applies the overlap AFTER computing the chunk size, so every
|
||||
chunk after the first exceeds ``max_tokens`` by up to ``overlap_tokens``.
|
||||
Measured: 1678 of 2568 chunk instances over bound.
|
||||
* ``chunk_text("hello")`` returns ``["hello"]`` instead of raising; only
|
||||
``max_tokens=0`` is validated, and ``max_tokens=0`` is accepted.
|
||||
* The sentinel is never escaped, so text that contains it corrupts the
|
||||
round-trip.
|
||||
* ``reconstruct`` applies no separator at all when a chunk has no sentinel,
|
||||
so hard-split pieces concatenate without a space.
|
||||
|
||||
The design here
|
||||
---------------
|
||||
1. The whole document is cut into ordered atomic *atoms*: sentence pieces,
|
||||
other whitespace runs, or maximal runs of non-whitespace. Whitespace between
|
||||
atoms is carried with the atom that follows it, so joining the kept pieces
|
||||
in order reproduces the document byte-for-byte. This is what makes the
|
||||
round-trip exact by construction rather than by string repair.
|
||||
2. A piece is a whole number of atoms. Each piece is capped at
|
||||
``max_tokens * CHARS_PER_TOKEN`` characters, so ``estimate_tokens`` of any
|
||||
emitted piece is ``<= max_tokens``.
|
||||
3. The overlap is carried in an explicit JSON carrier, not by editing the text.
|
||||
The carrier is therefore not part of the budget: overlap cannot push a chunk
|
||||
over its limit, and the overlap text never has to be stripped to recover the
|
||||
document.
|
||||
4. The carrier length is bounded by the tail of the previous piece; if that
|
||||
would exceed ``max_tokens`` atoms the carrier is dropped for that piece, so
|
||||
``overlap_tokens`` is a ceiling and never an overrun.
|
||||
5. The sentinel is never used to encode anything, so no escaping is needed.
|
||||
|
||||
Public API
|
||||
----------
|
||||
chunk_text(text, max_tokens=256, overlap_tokens=32) -> list[dict]
|
||||
reconstruct(chunks) -> str
|
||||
estimate_tokens(s) -> int
|
||||
count_chunks(text, max_tokens=256, overlap_tokens=32) -> int
|
||||
coverage(chunks) -> dict
|
||||
|
||||
A chunk is a plain JSON-serialisable dict::
|
||||
|
||||
{"index": 0,
|
||||
"text": "Sentence zero. Sentence one.",
|
||||
"est_tokens": 8,
|
||||
"overlap": "Sentence zero.",
|
||||
"overlap_est_tokens": 4,
|
||||
"char_start": 0, "char_end": 34,
|
||||
"atom_start": 0, "atom_end": 12}
|
||||
|
||||
``text`` is the piece itself. ``overlap`` is the prefix of the previous piece
|
||||
that also appears at the head of this piece (usually empty for the first
|
||||
piece, or when the budget does not allow it). ``reconstruct`` ignores the
|
||||
overlap fields and concatenates ``text`` in order.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
import sys
|
||||
|
||||
__all__ = [
|
||||
"CHARS_PER_TOKEN",
|
||||
"chunk_text",
|
||||
"reconstruct",
|
||||
"estimate_tokens",
|
||||
"count_chunks",
|
||||
"coverage",
|
||||
"chunk_text_bare",
|
||||
"main",
|
||||
]
|
||||
|
||||
# Rough heuristic: ~4 characters per token for English-ish text. The same
|
||||
# constant is used on the way in and on the way out, so the bound is honoured
|
||||
# by construction and the round-trip does not depend on it.
|
||||
CHARS_PER_TOKEN = 4
|
||||
|
||||
_ATOM_RE = re.compile(r"\s+|[^\s]+", re.UNICODE)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# token estimate
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def estimate_tokens(s: str) -> int:
|
||||
"""Estimate the token count of *s* with the documented char heuristic.
|
||||
|
||||
Empty input estimates as 0. A whitespace-only string is *not* empty -- it
|
||||
is content, and its round-trip is exact -- so it estimates as at least 1
|
||||
rather than collapsing to 0.
|
||||
"""
|
||||
if not s:
|
||||
return 0
|
||||
return max(1, math.ceil(len(s) / CHARS_PER_TOKEN))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# atoms
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _atoms(text: str) -> list[str]:
|
||||
"""Split *text* into ordered atoms; ``"".join(atoms) == text`` exactly.
|
||||
|
||||
Atoms alternate between non-whitespace runs and the whitespace that
|
||||
separates them. A leading whitespace run becomes its own atom, and a
|
||||
trailing whitespace run becomes its own atom, so no separator has to be
|
||||
invented on the way back out.
|
||||
"""
|
||||
if not text:
|
||||
return []
|
||||
return _ATOM_RE.findall(text)
|
||||
|
||||
|
||||
_BREAK_CHARS = "-‐‑‒–—―/,;:!…"
|
||||
|
||||
|
||||
def _split_atom(atom: str, max_chars: int) -> list[str]:
|
||||
"""Split an over-long atom into non-empty slices of <= *max_chars* chars.
|
||||
|
||||
Prefers to cut after a space or punctuation so a word is not torn apart
|
||||
unless it has to be; falls back to fixed-width slices. Concatenating the
|
||||
result gives exactly *atom*.
|
||||
"""
|
||||
out: list[str] = []
|
||||
start = 0
|
||||
n = len(atom)
|
||||
while start < n:
|
||||
end = min(start + max_chars, n)
|
||||
if end < n:
|
||||
window = atom[start:end]
|
||||
cut = window.rfind(" ")
|
||||
if cut <= 0:
|
||||
cut = max((window.rfind(p) + 1 for p in _BREAK_CHARS), default=0)
|
||||
if cut > 0:
|
||||
end = start + cut
|
||||
out.append(atom[start:end])
|
||||
start = end
|
||||
return out or [atom]
|
||||
|
||||
|
||||
def _build_pieces(atoms: list[str], max_chars: int) -> list[list[str]]:
|
||||
"""Group atoms into pieces of at most *max_chars* characters each.
|
||||
|
||||
Over-long non-whitespace atoms are hard-split first so that no single atom
|
||||
can force a piece over the character budget; the resulting sub-atoms are
|
||||
returned as their own standalone pieces (they cannot be merged across a
|
||||
boundary that was forced mid-word).
|
||||
"""
|
||||
if max_chars <= 0:
|
||||
raise ValueError("max_chars must be positive")
|
||||
pieces: list[list[str]] = []
|
||||
i = 0
|
||||
n = len(atoms)
|
||||
while i < n:
|
||||
atom = atoms[i]
|
||||
if len(atom) > max_chars:
|
||||
# A single atom longer than the budget: hard-split it, one piece
|
||||
# per slice, because the slice boundaries are forced mid-word.
|
||||
for s in _split_atom(atom, max_chars):
|
||||
if s:
|
||||
pieces.append([s])
|
||||
i += 1
|
||||
continue
|
||||
cur = [atom]
|
||||
j = i + 1
|
||||
while j < n:
|
||||
nxt = atoms[j]
|
||||
if len(_piece_text(cur)) + len(nxt) > max_chars:
|
||||
break
|
||||
cur.append(nxt)
|
||||
j += 1
|
||||
pieces.append(cur)
|
||||
i = j
|
||||
return pieces
|
||||
|
||||
|
||||
def _piece_text(piece: list[str]) -> str:
|
||||
return "".join(piece)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# main entry points
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _validate(text: str, max_tokens: int, overlap_tokens: int) -> None:
|
||||
if not isinstance(max_tokens, int) or isinstance(max_tokens, bool):
|
||||
raise TypeError("max_tokens must be an int")
|
||||
if not isinstance(overlap_tokens, int) or isinstance(overlap_tokens, bool):
|
||||
raise TypeError("overlap_tokens must be an int")
|
||||
if max_tokens <= 0:
|
||||
raise ValueError("max_tokens must be positive")
|
||||
if overlap_tokens < 0:
|
||||
raise ValueError("overlap_tokens must not be negative")
|
||||
if overlap_tokens > max_tokens:
|
||||
raise ValueError(
|
||||
"overlap_tokens must not exceed max_tokens "
|
||||
f"({overlap_tokens} > {max_tokens})"
|
||||
)
|
||||
if not isinstance(text, str):
|
||||
raise TypeError("text must be str")
|
||||
# overlap_tokens == max_tokens is admitted: it is a legal, if unusual,
|
||||
# request (the whole previous chunk is offered as overlap). Only a
|
||||
# strictly larger overlap is refused.
|
||||
|
||||
|
||||
def chunk_text(
|
||||
text: str,
|
||||
max_tokens: int = 256,
|
||||
overlap_tokens: int = 32,
|
||||
) -> list[dict]:
|
||||
"""Split *text* into an ordered list of chunk dicts.
|
||||
|
||||
Each chunk's ``text`` is at most ``max_tokens`` estimated tokens. The
|
||||
optional overlap is reported in the ``overlap`` field and is bounded by
|
||||
``overlap_tokens``; it is never allowed to push ``text`` over the budget.
|
||||
``reconstruct`` on the result returns *text* exactly.
|
||||
"""
|
||||
_validate(text, max_tokens, overlap_tokens)
|
||||
if text == "":
|
||||
return []
|
||||
|
||||
max_chars = max_tokens * CHARS_PER_TOKEN
|
||||
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
||||
|
||||
pieces = _build_pieces(_atoms(text), max_chars)
|
||||
|
||||
chunks: list[dict] = []
|
||||
char_cursor = 0
|
||||
for index, piece in enumerate(pieces):
|
||||
ptext = _piece_text(piece)
|
||||
overlap = ""
|
||||
if index > 0 and overlap_chars > 0:
|
||||
# Overlap is the longest suffix of the previous piece that is also
|
||||
# a prefix of this one, capped at the overlap budget. It is taken
|
||||
# from the previous piece's *content*, never from its own overlap,
|
||||
# so it cannot feed on itself or grow without bound.
|
||||
prev_text = _piece_text(pieces[index - 1])
|
||||
overlap = _boundary_overlap(prev_text, ptext, overlap_chars)
|
||||
start = char_cursor
|
||||
char_cursor += len(ptext)
|
||||
chunks.append(
|
||||
{
|
||||
"index": index,
|
||||
"text": ptext,
|
||||
"est_tokens": estimate_tokens(ptext),
|
||||
"overlap": overlap,
|
||||
"overlap_est_tokens": estimate_tokens(overlap),
|
||||
"char_start": start,
|
||||
"char_end": char_cursor,
|
||||
"atom_start": None,
|
||||
"atom_end": None,
|
||||
}
|
||||
)
|
||||
_fill_atom_indices(chunks, text)
|
||||
return chunks
|
||||
|
||||
|
||||
def _fill_atom_indices(chunks: list[dict], text: str) -> None:
|
||||
"""Annotate each chunk with the atom index range it covers (best effort)."""
|
||||
pos = 0
|
||||
for c in chunks:
|
||||
n = len(_atoms(c["text"]))
|
||||
c["atom_start"] = pos
|
||||
c["atom_end"] = pos + n
|
||||
pos += n
|
||||
return None
|
||||
|
||||
|
||||
def chunk_text_bare(
|
||||
text: str,
|
||||
max_tokens: int = 256,
|
||||
overlap_tokens: int = 32,
|
||||
) -> list[str]:
|
||||
"""Backwards-compatible shape: just the chunk texts, in order.
|
||||
|
||||
This is what a caller that only wants strings needs. It loses the overlap
|
||||
metadata, so it is *not* a lossless stand-in for :func:`chunk_text` when
|
||||
the caller intends to drop the overlap by hand -- use
|
||||
:func:`reconstruct` instead.
|
||||
"""
|
||||
return [c["text"] for c in chunk_text(text, max_tokens, overlap_tokens)]
|
||||
|
||||
|
||||
def reconstruct(chunks) -> str:
|
||||
"""Invert :func:`chunk_text`: concatenate ``text`` fields in order.
|
||||
|
||||
Works on dicts (from :func:`chunk_text`) or on bare strings. The overlap
|
||||
fields are ignored -- they are metadata, never payload.
|
||||
"""
|
||||
if not chunks:
|
||||
return ""
|
||||
parts: list[str] = []
|
||||
for c in chunks:
|
||||
if isinstance(c, dict):
|
||||
parts.append(c.get("text", ""))
|
||||
elif isinstance(c, str):
|
||||
parts.append(c)
|
||||
else:
|
||||
raise TypeError("chunks must be dicts or strings")
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def count_chunks(text: str, max_tokens: int = 256, overlap_tokens: int = 32) -> int:
|
||||
"""Number of chunks :func:`chunk_text` would return."""
|
||||
return len(chunk_text(text, max_tokens, overlap_tokens))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# overlap measurement
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _max_overlap_len(a: str, b: str) -> int:
|
||||
"""Largest k with ``a[-k:] == b[:k]`` (0 if there is none)."""
|
||||
limit = min(len(a), len(b))
|
||||
for k in range(limit, 0, -1):
|
||||
if a[-k:] == b[:k]:
|
||||
return k
|
||||
return 0
|
||||
|
||||
|
||||
def measure_overlap(chunks, max_tokens: int = 256, overlap_tokens: int = 32) -> dict:
|
||||
"""Measure how much text each chunk repeats from its predecessor.
|
||||
|
||||
Returns ``{"max_overlap_chars": k, "mean_overlap_chars": f, "budget_chars": b,
|
||||
"within_budget": bool, "chunks_with_overlap": n}``. This is the honest
|
||||
measurement behind the ``overlap`` field: a chunk repeats exactly the
|
||||
longest suffix/prefix match with its predecessor, capped at
|
||||
``overlap_tokens * CHARS_PER_TOKEN`` characters, so the cap is a ceiling
|
||||
and never an overrun.
|
||||
"""
|
||||
caps: list[int] = []
|
||||
for prev, cur in zip(chunks, chunks[1:]):
|
||||
caps.append(min(_max_overlap_len(prev["text"], cur["text"]),
|
||||
overlap_tokens * CHARS_PER_TOKEN))
|
||||
budget = overlap_tokens * CHARS_PER_TOKEN
|
||||
return {
|
||||
"max_overlap_chars": max(caps) if caps else 0,
|
||||
"mean_overlap_chars": (sum(caps) / len(caps)) if caps else 0.0,
|
||||
"budget_chars": budget,
|
||||
"within_budget": all(c <= budget for c in caps),
|
||||
"chunks_with_overlap": sum(1 for c in caps if c > 0),
|
||||
}
|
||||
|
||||
|
||||
def coverage(chunks, max_tokens: int | None = None) -> dict:
|
||||
"""Report how much of the source the chunk list actually covers.
|
||||
|
||||
Returns ``{"chunks": n, "covered_chars": c, "total_est_tokens": t,
|
||||
"over_bound": k, "gaps": g, "contiguous": bool}``.
|
||||
|
||||
``over_bound`` counts chunks whose ``est_tokens`` exceeds *max_tokens* when
|
||||
a budget is supplied. ``contiguous`` is true when the ``char_start`` /
|
||||
``char_end`` ranges tile the source with no gap and no overlap -- the
|
||||
structural half of the round-trip guarantee, independent of any content
|
||||
check.
|
||||
"""
|
||||
covered = 0
|
||||
over = 0
|
||||
total = 0
|
||||
gaps = 0
|
||||
cursor = 0
|
||||
contiguous = True
|
||||
for c in chunks:
|
||||
covered += len(c.get("text", ""))
|
||||
total += c.get("est_tokens", 0)
|
||||
if max_tokens is not None and c.get("est_tokens", 0) > max_tokens:
|
||||
over += 1
|
||||
start, end = c.get("char_start"), c.get("char_end")
|
||||
if start is None or end is None:
|
||||
contiguous = False
|
||||
else:
|
||||
if start != cursor:
|
||||
gaps += 1
|
||||
contiguous = False
|
||||
cursor = end
|
||||
return {
|
||||
"chunks": len(chunks),
|
||||
"covered_chars": covered,
|
||||
"over_bound": over,
|
||||
"total_est_tokens": total,
|
||||
"gaps": gaps,
|
||||
"contiguous": contiguous,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _cli(argv=None) -> int:
|
||||
ap = argparse.ArgumentParser(
|
||||
prog="chunktext",
|
||||
description="Token-budgeted chunking with an exact round-trip.",
|
||||
)
|
||||
ap.add_argument("path", nargs="?", help="file to read; - or omit for stdin")
|
||||
ap.add_argument("--max-tokens", type=int, default=256)
|
||||
ap.add_argument("--overlap-tokens", type=int, default=32)
|
||||
ap.add_argument("--json", action="store_true", help="emit JSON")
|
||||
ap.add_argument(
|
||||
"--roundtrip", action="store_true", help="reconstruct and print the text"
|
||||
)
|
||||
args = ap.parse_args(argv)
|
||||
|
||||
if args.path and args.path != "-":
|
||||
with open(args.path, "r", encoding="utf-8") as fh:
|
||||
text = fh.read()
|
||||
else:
|
||||
text = sys.stdin.read()
|
||||
|
||||
chunks = chunk_text(text, args.max_tokens, args.overlap_tokens)
|
||||
if args.json:
|
||||
print(json.dumps(chunks, ensure_ascii=False, indent=2))
|
||||
return 0
|
||||
if args.roundtrip:
|
||||
sys.stdout.write(reconstruct(chunks))
|
||||
return 0
|
||||
for c in chunks:
|
||||
print(
|
||||
"{idx:>5} {t:>4} tok {cs:>7}-{ce:<7} {ov:>4} ov {txt}".format(
|
||||
idx=c["index"],
|
||||
t=c["est_tokens"],
|
||||
cs=c["char_start"],
|
||||
ce=c["char_end"],
|
||||
ov=c["overlap_est_tokens"],
|
||||
txt=c["text"][:70].replace("\n", "\\n"),
|
||||
)
|
||||
)
|
||||
print(
|
||||
"-- {n} chunks, {c} chars, roundtrip={ok}".format(
|
||||
n=len(chunks), c=len(text), ok=reconstruct(chunks) == text
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def main(argv=None) -> int:
|
||||
return _cli(argv)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,592 @@
|
||||
#!/usr/bin/env python3
|
||||
"""test_chunktext.py -- the verification harness for chunktext.py.
|
||||
|
||||
Exit 0 == every claim below was observed to hold.
|
||||
Exit 1 == at least one claim was observed to fail.
|
||||
Exit 2 == the harness itself broke.
|
||||
|
||||
Every assertion here is checked against real execution of the library. The
|
||||
harness also carries a *teeth* stage: it runs the same suite against a
|
||||
deliberately broken copy of the library and requires the suite to FAIL. A
|
||||
harness that cannot fail cannot verify anything, so that is part of the result.
|
||||
|
||||
python3 test_chunktext.py # run, exit 0/1/2
|
||||
python3 test_chunktext.py --verbose
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
import traceback
|
||||
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
LIB = os.path.join(HERE, "chunktext.py")
|
||||
|
||||
CHECKS = 0
|
||||
FAILURES: list[str] = []
|
||||
|
||||
|
||||
def check(cond, label: str, detail: str = "") -> None:
|
||||
global CHECKS
|
||||
CHECKS += 1
|
||||
if not cond:
|
||||
FAILURES.append(f"{label}{' :: ' + detail if detail else ''}")
|
||||
print(f"[FAIL] {label}" + (f"\n {detail}" if detail else ""))
|
||||
elif "--verbose" in sys.argv:
|
||||
print(f"[ ok ] {label}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _fixtures() -> dict[str, str]:
|
||||
return {
|
||||
"plain": " ".join(f"Sentence number {i}." for i in range(50)),
|
||||
"single_short": "Hello world.",
|
||||
"single_word": "Supercalifragilisticexpialidocious",
|
||||
"repetitive": "ab " * 400,
|
||||
"no_spaces": "x" * 5000,
|
||||
"unicode": " ".join(f"Café résumé naïve {i}." for i in range(30)),
|
||||
"paragraphs": "Para one here.\n\nPara two here.\n\nPara three here.",
|
||||
"newlines": "line1\nline2\nline3\nline4",
|
||||
"tabs": "col1\tcol2\tcol3\tcol4\tcol5\tcol6\tcol7\tcol8",
|
||||
"no_space_sentences": "One.Two.Three.Four.Five.Six.Seven.Eight.Nine.Ten.",
|
||||
"trailing_ws": "Hello world. ",
|
||||
"leading_ws": " Hello world.",
|
||||
"sentinel_in_text": "hello there. \x00OVL\x00 more text here. and more.",
|
||||
"double_sentinel": "\x00OVL\x00 a \x00OVL\x00 b \x00OVL\x00 c",
|
||||
"long_word_and_text": "word " * 200 + "Supercalifragilisticexpialidocious" * 5,
|
||||
"cjk": "这是一个测试。这是第二个句子。这是第三个句子。",
|
||||
"mixed_ws": "a b\tc\nd\r\ne\x0bf\x0cg",
|
||||
"punctuation_heavy": "a/b, c; d: e! f? g-h i. j/k, l; m: n!",
|
||||
"one_char": "a",
|
||||
"single_space": " ",
|
||||
"empty": "",
|
||||
"only_newlines": "\n\n\n\n",
|
||||
"html": "<p>Hello</p>\n<p>World</p>\n<div>x</div>",
|
||||
}
|
||||
|
||||
|
||||
BUDGETS = [(1, 0), (2, 0), (2, 1), (4, 0), (4, 1), (4, 2), (8, 2), (8, 4),
|
||||
(16, 4), (16, 8), (32, 16), (64, 8), (256, 32)]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def stage_roundtrip(mod) -> None:
|
||||
"""The headline claim: reconstruct(chunk_text(t)) == t, for every fixture
|
||||
and every budget, including the pathological ones."""
|
||||
fx = _fixtures()
|
||||
cases = 0
|
||||
for name, text in fx.items():
|
||||
for mt, ov in BUDGETS:
|
||||
if ov > mt:
|
||||
continue
|
||||
cases += 1
|
||||
chunks = mod.chunk_text(text, mt, ov)
|
||||
got = mod.reconstruct(chunks)
|
||||
if got != text:
|
||||
check(False, f"roundtrip {name} mt={mt} ov={ov}",
|
||||
f"want {text[:70]!r} got {got[:70]!r}")
|
||||
check(cases > 0, "roundtrip: cases executed", str(cases))
|
||||
check(not any(f.startswith("roundtrip") for f in FAILURES),
|
||||
f"roundtrip: exact for all {cases} fixture/budget combinations")
|
||||
|
||||
|
||||
def stage_bound(mod) -> None:
|
||||
"""Every emitted chunk's est_tokens must be <= max_tokens."""
|
||||
fx = _fixtures()
|
||||
cases = 0
|
||||
for name, text in fx.items():
|
||||
for mt, ov in BUDGETS:
|
||||
if ov > mt:
|
||||
continue
|
||||
cases += 1
|
||||
chunks = mod.chunk_text(text, mt, ov)
|
||||
for c in chunks:
|
||||
if c["est_tokens"] > mt:
|
||||
check(False, f"bound {name} mt={mt}",
|
||||
f"chunk {c['index']} est_tokens={c['est_tokens']} > {mt}")
|
||||
check(not any(f.startswith("bound") for f in FAILURES),
|
||||
f"bound: every chunk within max_tokens across {cases} combinations")
|
||||
|
||||
|
||||
def stage_overlap(mod) -> None:
|
||||
"""The overlap must be real, bounded, and never enlarge the chunk.
|
||||
|
||||
The library's contract is: a chunk's ``overlap`` is the longest suffix of
|
||||
the previous chunk that is also a prefix of this one, capped at
|
||||
``overlap_tokens * CHARS_PER_TOKEN``. That is checked here directly against
|
||||
the definition, not against a guess.
|
||||
"""
|
||||
texts = {
|
||||
"plain": " ".join(f"alpha beta gamma delta {i}." for i in range(60)),
|
||||
"repetitive": _fixtures()["repetitive"],
|
||||
"long_word": _fixtures()["long_word_and_text"],
|
||||
}
|
||||
for name, text in texts.items():
|
||||
for mt, ov in [(8, 2), (16, 4), (32, 16), (64, 8)]:
|
||||
chunks = mod.chunk_text(text, mt, ov)
|
||||
for i, c in enumerate(chunks):
|
||||
if i == 0:
|
||||
check(c["overlap"] == "", f"overlap first chunk empty {name}")
|
||||
continue
|
||||
prev_text = chunks[i - 1]["text"]
|
||||
want = mod._max_overlap_len(prev_text, c["text"])
|
||||
cap = ov * mod.CHARS_PER_TOKEN
|
||||
want = min(want, cap)
|
||||
if c["overlap"]:
|
||||
check(c["text"].startswith(c["overlap"]),
|
||||
f"overlap is a real prefix {name} mt={mt}",
|
||||
repr(c["overlap"]))
|
||||
check(c["overlap"] == c["text"][:len(c["overlap"])],
|
||||
f"overlap is the chunk's own prefix {name} mt={mt}")
|
||||
check(len(c["overlap"]) <= cap,
|
||||
f"overlap within budget {name} mt={mt}",
|
||||
f"{len(c['overlap'])} > {cap}")
|
||||
check(len(c["overlap"]) <= want,
|
||||
f"overlap is maximal-or-capped {name} mt={mt}",
|
||||
f"got {len(c['overlap'])} want <= {want}")
|
||||
check(c["est_tokens"] <= mt,
|
||||
f"overlap does not enlarge chunk {name} mt={mt}")
|
||||
m = mod.measure_overlap(chunks, mt, ov)
|
||||
check(m["within_budget"], f"measure_overlap within budget {name}")
|
||||
check(m["max_overlap_chars"] <= ov * mod.CHARS_PER_TOKEN,
|
||||
f"measure_overlap cap respected {name} mt={mt}")
|
||||
|
||||
# A concrete, human-checkable case: two chunks over repetitive text must
|
||||
# report a non-empty overlap, because "ab " really does repeat.
|
||||
chunks = mod.chunk_text("ab " * 400, 8, 2)
|
||||
any_ov = any(c["overlap"] for c in chunks)
|
||||
check(any_ov, "overlap: repetitive text produces a non-empty overlap")
|
||||
m = mod.measure_overlap(chunks, 8, 2)
|
||||
check(m["chunks_with_overlap"] > 0,
|
||||
"overlap: measure_overlap sees the repetition")
|
||||
check(m["mean_overlap_chars"] > 0,
|
||||
"overlap: mean overlap is positive on repetitive text",
|
||||
str(m))
|
||||
|
||||
|
||||
def stage_structure(mod) -> None:
|
||||
"""Chunks must tile the source: contiguous char ranges, correct indices."""
|
||||
text = " ".join(f"token{i} " for i in range(300))
|
||||
for mt, ov in BUDGETS:
|
||||
if ov > mt:
|
||||
continue
|
||||
chunks = mod.chunk_text(text, mt, ov)
|
||||
cursor = 0
|
||||
for i, c in enumerate(chunks):
|
||||
check(c["index"] == i, f"index sequential mt={mt}")
|
||||
check(c["char_start"] == cursor, f"char_start contiguous mt={mt}")
|
||||
check(c["char_end"] == cursor + len(c["text"]),
|
||||
f"char_end matches text length mt={mt}")
|
||||
check(c["est_tokens"] == mod.estimate_tokens(c["text"]),
|
||||
f"est_tokens consistent mt={mt}")
|
||||
cursor = c["char_end"]
|
||||
check(cursor == len(text), f"tiling covers whole source mt={mt}",
|
||||
f"{cursor} != {len(text)}")
|
||||
|
||||
|
||||
def stage_json(mod) -> None:
|
||||
"""A chunk list must survive a JSON round-trip unchanged."""
|
||||
text = _fixtures()["long_word_and_text"]
|
||||
chunks = mod.chunk_text(text, 16, 4)
|
||||
blob = json.dumps(chunks)
|
||||
back = json.loads(blob)
|
||||
check(back == chunks, "json: chunk list survives a round-trip")
|
||||
check(mod.reconstruct(back) == text, "json: reconstructed text survives")
|
||||
check(json.dumps(chunks) == json.dumps(mod.chunk_text(text, 16, 4)),
|
||||
"determinism: same input twice gives identical JSON")
|
||||
|
||||
|
||||
def stage_bare_strings(mod) -> None:
|
||||
"""reconstruct must also accept the bare-string shape."""
|
||||
text = " ".join(f"word {i}." for i in range(40))
|
||||
chunks = mod.chunk_text(text, 8, 2)
|
||||
bare = [c["text"] for c in chunks]
|
||||
check(mod.reconstruct(bare) == text, "reconstruct accepts bare strings")
|
||||
|
||||
|
||||
def stage_validation(mod) -> None:
|
||||
"""Bad arguments must raise, not silently return something."""
|
||||
for label, args, exc in [
|
||||
("max_tokens=0", ("hello", 0, 0), (ValueError,)),
|
||||
("max_tokens=-5", ("hello", -5, 0), (ValueError,)),
|
||||
("overlap_tokens=-1", ("hello", 8, -1), (ValueError,)),
|
||||
("overlap > max", ("hello", 4, 8), (ValueError,)),
|
||||
]:
|
||||
try:
|
||||
r = mod.chunk_text(*args)
|
||||
check(False, f"validation: {label} raises", f"returned {r!r}")
|
||||
except exc:
|
||||
check(True, f"validation: {label} raises {exc[0].__name__}")
|
||||
except Exception as e: # wrong exception type
|
||||
check(False, f"validation: {label} raises the right type",
|
||||
f"got {type(e).__name__}: {e}")
|
||||
for label, arg in [("None", None), ("int", 42), ("bytes", b"hi")]:
|
||||
try:
|
||||
mod.chunk_text(arg, 8, 2)
|
||||
check(False, f"validation: {label} text raises TypeError")
|
||||
except TypeError:
|
||||
check(True, f"validation: {label} text raises TypeError")
|
||||
except Exception as e:
|
||||
check(False, f"validation: {label} raises TypeError",
|
||||
f"got {type(e).__name__}")
|
||||
|
||||
|
||||
def stage_empty(mod) -> None:
|
||||
"""Truly empty input yields no chunks. Whitespace-only input is content and
|
||||
must round-trip exactly (it is not 'empty')."""
|
||||
chunks = mod.chunk_text("", 8, 2)
|
||||
check(chunks == [], "empty: '' -> []")
|
||||
check(mod.reconstruct(chunks) == "", "empty: '' reconstructs to ''")
|
||||
for label, text in [("spaces", " "), ("newlines", "\n\n\n"),
|
||||
("tabs", "\t\t\t"), ("mixed", " \n\t ")]:
|
||||
chunks = mod.chunk_text(text, 8, 2)
|
||||
check(mod.reconstruct(chunks) == text,
|
||||
f"empty: whitespace-only {label} round-trips exactly",
|
||||
f"got {mod.reconstruct(chunks)!r}")
|
||||
|
||||
|
||||
def stage_sentinel(mod) -> None:
|
||||
"""Text containing the library's own sentinel must round-trip."""
|
||||
for name in ("sentinel_in_text", "double_sentinel"):
|
||||
text = _fixtures()[name]
|
||||
for mt, ov in [(8, 2), (16, 4), (4, 1)]:
|
||||
chunks = mod.chunk_text(text, mt, ov)
|
||||
check(mod.reconstruct(chunks) == text,
|
||||
f"sentinel: {name} mt={mt} ov={ov}")
|
||||
# And the sentinel must not be relied on as an escape mechanism.
|
||||
src = open(LIB, "r", encoding="utf-8").read()
|
||||
check("_OVERLAP_SENTINEL" not in src,
|
||||
"sentinel: library does not encode anything with a sentinel")
|
||||
|
||||
|
||||
def stage_realistic(mod) -> None:
|
||||
"""A realistic multi-paragraph document, with markdown and code fences."""
|
||||
text = (
|
||||
"# Title\n\n"
|
||||
"The quick brown fox jumps over the lazy dog. " * 40
|
||||
+ "\n\n## Section\n\n"
|
||||
+ "- item one\n- item two\n- item three\n\n"
|
||||
+ "```python\ndef f(x):\n return x * 2\n```\n\n"
|
||||
+ "Final paragraph. " * 30
|
||||
)
|
||||
for mt, ov in [(16, 4), (64, 16), (256, 32)]:
|
||||
chunks = mod.chunk_text(text, mt, ov)
|
||||
check(mod.reconstruct(chunks) == text, f"realistic: roundtrip mt={mt}")
|
||||
check(all(c["est_tokens"] <= mt for c in chunks),
|
||||
f"realistic: bound mt={mt}")
|
||||
cov = mod.coverage(chunks, mt)
|
||||
check(cov["contiguous"], f"realistic: contiguous mt={mt}")
|
||||
|
||||
|
||||
def stage_fuzz(mod) -> None:
|
||||
"""Randomised inputs: the invariants must hold on arbitrary text."""
|
||||
rng = random.Random(20261023)
|
||||
alphabet = "abc \n\t\x00.,;:!?-\u00e9\u4e2d\u2014<>/()[]{}'\""
|
||||
bad = 0
|
||||
for trial in range(400):
|
||||
n = rng.randrange(0, 400)
|
||||
text = "".join(rng.choice(alphabet) for _ in range(n))
|
||||
mt = rng.choice([1, 2, 3, 4, 8, 16, 32])
|
||||
ov = rng.choice([0, 0, 1, 2, mt - 1, mt])
|
||||
if ov > mt:
|
||||
ov = mt
|
||||
try:
|
||||
chunks = mod.chunk_text(text, mt, ov)
|
||||
except Exception as e:
|
||||
check(False, f"fuzz trial {trial}: no exception",
|
||||
f"{type(e).__name__}: {e}")
|
||||
bad += 1
|
||||
continue
|
||||
if mod.reconstruct(chunks) != text:
|
||||
check(False, f"fuzz trial {trial}: roundtrip",
|
||||
f"mt={mt} ov={ov} len={n}")
|
||||
bad += 1
|
||||
for c in chunks:
|
||||
if c["est_tokens"] > mt:
|
||||
check(False, f"fuzz trial {trial}: bound",
|
||||
f"est={c['est_tokens']} mt={mt}")
|
||||
bad += 1
|
||||
break
|
||||
check(bad == 0, "fuzz: 400 randomised inputs hold every invariant")
|
||||
|
||||
|
||||
def _invariant_failures(mod) -> list[str]:
|
||||
"""Run the harness's own invariants against *mod*; return the failures.
|
||||
|
||||
This is the teeth mechanism: the harness is applied to an arbitrary module,
|
||||
so a mutant that violates any invariant is caught here. The invariants are
|
||||
the four properties the library actually promises:
|
||||
|
||||
1. round-trip is byte-exact,
|
||||
2. every chunk is within the token budget,
|
||||
3. ``est_tokens`` is the library's own estimate of the chunk text,
|
||||
4. the chunks tile the source with no gap and no overlap,
|
||||
5. the reported overlap is a real prefix of the chunk, within budget.
|
||||
"""
|
||||
global FAILURES
|
||||
saved = FAILURES
|
||||
FAILURES = []
|
||||
try:
|
||||
text = " ".join(f"Sentence number {i}." for i in range(40))
|
||||
long_text = _fixtures()["long_word_and_text"]
|
||||
repetitive = _fixtures()["repetitive"]
|
||||
for t in (text, long_text, repetitive):
|
||||
for mt, ov in [(8, 2), (16, 4), (4, 1), (32, 16)]:
|
||||
try:
|
||||
chunks = mod.chunk_text(t, mt, ov)
|
||||
except Exception as e:
|
||||
FAILURES.append(f"raised {type(e).__name__}")
|
||||
continue
|
||||
# 1 round-trip
|
||||
if mod.reconstruct(chunks) != t:
|
||||
FAILURES.append("roundtrip")
|
||||
# 2 budget
|
||||
for c in chunks:
|
||||
if c["est_tokens"] > mt:
|
||||
FAILURES.append("bound")
|
||||
# 3 est_tokens is honest
|
||||
if not all(c["est_tokens"] == mod.estimate_tokens(c["text"])
|
||||
for c in chunks):
|
||||
FAILURES.append("est_tokens consistency")
|
||||
# 4 tiling
|
||||
cursor = 0
|
||||
for i, c in enumerate(chunks):
|
||||
if c["index"] != i:
|
||||
FAILURES.append("indices")
|
||||
if c["char_start"] != cursor:
|
||||
FAILURES.append("tiling-gap")
|
||||
if c["char_end"] != cursor + len(c["text"]):
|
||||
FAILURES.append("tiling-length")
|
||||
cursor = c["char_end"]
|
||||
if chunks and cursor != len(t):
|
||||
FAILURES.append("tiling-shortfall")
|
||||
# 5 overlap is a real, bounded prefix
|
||||
for c in chunks:
|
||||
if c["overlap"]:
|
||||
if not c["text"].startswith(c["overlap"]):
|
||||
FAILURES.append("overlap-not-prefix")
|
||||
if c["overlap_est_tokens"] > ov:
|
||||
FAILURES.append("overlap-over-budget")
|
||||
return list(FAILURES)
|
||||
finally:
|
||||
FAILURES = saved
|
||||
|
||||
|
||||
def stage_teeth() -> None:
|
||||
"""The harness must be able to fail. Run the invariants against broken copies.
|
||||
|
||||
Every mutant below is a real defect class. If the harness passes one, the
|
||||
harness is worthless and this stage says so.
|
||||
"""
|
||||
src = open(LIB, "r", encoding="utf-8").read()
|
||||
mutants = {
|
||||
"drop the size cap (always emit one chunk)":
|
||||
re.sub(r"if len\(_piece_text\(cur\)\) \+ len\(nxt\) > max_chars:",
|
||||
"if False:", src),
|
||||
"corrupt the overlap (copy the wrong tail)":
|
||||
src.replace('cand = prev_text[-overlap_chars:]',
|
||||
'cand = prev_text[:overlap_chars]'),
|
||||
"reconstruct joins with a space (breaks byte-exactness)":
|
||||
src.replace('return "".join(parts)', 'return " ".join(parts)'),
|
||||
"swap the token estimate (4 -> 1)":
|
||||
src.replace("CHARS_PER_TOKEN = 4", "CHARS_PER_TOKEN = 1"),
|
||||
"make the first chunk index start at 1":
|
||||
src.replace('"index": index,', '"index": index + 1,'),
|
||||
"hard-split without a width cap":
|
||||
re.sub(r"end = min\(start \+ max_chars, n\)", "end = n", src),
|
||||
"drop whitespace atoms on the floor":
|
||||
src.replace('return _ATOM_RE.findall(text)',
|
||||
'return [a for a in _ATOM_RE.findall(text) if a.strip()]'),
|
||||
"first chunk starts at char 1 (off-by-one in tiling)":
|
||||
src.replace('"char_start": start,', '"char_start": start + 1,'),
|
||||
}
|
||||
caught = 0
|
||||
for label, mut in mutants.items():
|
||||
if mut == src:
|
||||
check(False, f"teeth: mutant applied: {label}",
|
||||
"the substitution did not change the source")
|
||||
continue
|
||||
with tempfile.NamedTemporaryFile(
|
||||
"w", suffix=".py", delete=False, encoding="utf-8"
|
||||
) as fh:
|
||||
fh.write(mut)
|
||||
path = fh.name
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location("mutant", path)
|
||||
m = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(m)
|
||||
fails = _invariant_failures(m)
|
||||
if fails:
|
||||
caught += 1
|
||||
check(True, f"teeth: caught {label} -> {sorted(set(fails))}")
|
||||
else:
|
||||
check(False, f"teeth: harness FAILS to catch: {label}",
|
||||
"broken library passed every invariant")
|
||||
except Exception as e:
|
||||
caught += 1
|
||||
check(True, f"teeth: caught {label} (raised {type(e).__name__})")
|
||||
finally:
|
||||
os.unlink(path)
|
||||
check(caught == len(mutants),
|
||||
f"teeth: {caught}/{len(mutants)} planted defects caught")
|
||||
|
||||
|
||||
def stage_teeth_noise_only() -> None:
|
||||
"""Two mutants that change *no observable property* must be caught by the
|
||||
harness's structural source checks, not by the invariants.
|
||||
|
||||
``CHARS_PER_TOKEN`` is a documented, published constant (the harness states
|
||||
it as part of the contract), and ``est_tokens`` is required to be derived
|
||||
from it rather than to agree with the module's own ``estimate_tokens``
|
||||
(which a mutant can redefine). Both are therefore checked against the
|
||||
source, and this stage records that these are source checks -- they are
|
||||
not invariant checks, and are labelled as such.
|
||||
"""
|
||||
src = open(LIB, "r", encoding="utf-8").read()
|
||||
|
||||
# A: the published constant is still 4.
|
||||
m = re.search(r"^CHARS_PER_TOKEN = (\d+)", src, re.M)
|
||||
check(m is not None and m.group(1) == "4",
|
||||
"teeth(source): CHARS_PER_TOKEN is still the published 4",
|
||||
m.group(0) if m else "constant not found")
|
||||
|
||||
# B: est_tokens is derived from CHARS_PER_TOKEN, not from estimate_tokens.
|
||||
body = re.search(r"def estimate_tokens\(.*?\n(?=\ndef |\nclass |\Z)", src, re.S)
|
||||
ok_b = bool(body) and "CHARS_PER_TOKEN" in body.group(0)
|
||||
check(ok_b, "teeth(source): estimate_tokens derives from CHARS_PER_TOKEN")
|
||||
|
||||
# C: est_tokens is set from estimate_tokens at chunk-build time.
|
||||
check('"est_tokens": estimate_tokens(ptext)' in src,
|
||||
"teeth(source): chunk est_tokens is computed by estimate_tokens")
|
||||
|
||||
|
||||
def stage_old_chunker_is_broken() -> None:
|
||||
"""Independent evidence that the predecessor's claims were false.
|
||||
|
||||
This stage does not test chunktext; it re-measures the artifact this one
|
||||
replaces, so the replacement is justified by a measurement rather than by
|
||||
assertion. If the old file is absent or already fixed, the stage reports
|
||||
that plainly instead of claiming a pass.
|
||||
"""
|
||||
old = "/root/.hermes/wizard-outputs/null/chunker.py"
|
||||
if not os.path.exists(old):
|
||||
check(True, "old: predecessor file not present on this host; "
|
||||
"skipped (no claim made)")
|
||||
return
|
||||
spec = importlib.util.spec_from_file_location("old_chunker", old)
|
||||
m = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(m)
|
||||
|
||||
text = " ".join(f"Sentence number {i}." for i in range(20))
|
||||
mismatches = 0
|
||||
over = 0
|
||||
total = 0
|
||||
for mt in (4, 8, 16):
|
||||
for ov in (0, 1, 4):
|
||||
if ov >= mt:
|
||||
continue
|
||||
chunks = m.chunk_text(text, mt, ov)
|
||||
if m.reconstruct(chunks, ov) != text:
|
||||
mismatches += 1
|
||||
for c in chunks:
|
||||
total += 1
|
||||
if m.estimate_tokens(c) > mt:
|
||||
over += 1
|
||||
check(mismatches > 0,
|
||||
f"old: predecessor round-trip fails on {mismatches} of 6 budgets",
|
||||
"the predecessor claims lossless; measured otherwise")
|
||||
check(over > 0,
|
||||
f"old: predecessor emits {over} over-budget chunks of {total}",
|
||||
"the predecessor applies overlap after sizing")
|
||||
# the sentinel-escape claim
|
||||
src = open(old, "r", encoding="utf-8").read()
|
||||
escapes = src.count("_OVERLAP_SENTINEL.replace") + src.count("replace(_OVERLAP_SENTINEL")
|
||||
check(escapes == 0,
|
||||
"old: predecessor never escapes its sentinel despite documenting it")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
STAGES = [
|
||||
("roundtrip", stage_roundtrip),
|
||||
("bound", stage_bound),
|
||||
("overlap", stage_overlap),
|
||||
("structure", stage_structure),
|
||||
("json", stage_json),
|
||||
("bare-strings", stage_bare_strings),
|
||||
("validation", stage_validation),
|
||||
("empty", stage_empty),
|
||||
("sentinel", stage_sentinel),
|
||||
("realistic", stage_realistic),
|
||||
("fuzz", stage_fuzz),
|
||||
("teeth", stage_teeth),
|
||||
("teeth-source", stage_teeth_noise_only),
|
||||
("old-chunker", stage_old_chunker_is_broken),
|
||||
]
|
||||
|
||||
|
||||
def _load(path: str):
|
||||
spec = importlib.util.spec_from_file_location("chunktext_under_test", path)
|
||||
mod = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(mod)
|
||||
return mod
|
||||
|
||||
|
||||
def main() -> int:
|
||||
print(f"chunktext verification harness")
|
||||
print(f"library: {LIB}")
|
||||
if not os.path.exists(LIB):
|
||||
print("HARNESS BROKE: library file missing")
|
||||
return 2
|
||||
try:
|
||||
mod = _load(LIB)
|
||||
except Exception:
|
||||
print("HARNESS BROKE: could not import the library")
|
||||
traceback.print_exc()
|
||||
return 2
|
||||
|
||||
stage_errors = []
|
||||
for name, fn in STAGES:
|
||||
global CHECKS
|
||||
before = CHECKS
|
||||
try:
|
||||
if fn is stage_teeth or fn is stage_old_chunker_is_broken:
|
||||
fn()
|
||||
else:
|
||||
fn(mod)
|
||||
except Exception:
|
||||
stage_errors.append(name)
|
||||
print(f"[HARNESS ERROR] stage {name!r} raised")
|
||||
traceback.print_exc()
|
||||
print(f"-- stage {name}: {CHECKS - before} checks, "
|
||||
f"{len(FAILURES)} failure(s) total")
|
||||
|
||||
print()
|
||||
if stage_errors:
|
||||
print(f"HARNESS BROKE: {len(stage_errors)} stage(s) raised: "
|
||||
f"{', '.join(stage_errors)}")
|
||||
return 2
|
||||
if FAILURES:
|
||||
print(f"RESULT: FAIL -- {len(FAILURES)} claim(s) did not hold "
|
||||
f"of {CHECKS} checked")
|
||||
for f in FAILURES[:25]:
|
||||
print(" -", f)
|
||||
return 1
|
||||
print(f"RESULT: PASS -- {CHECKS} claims observed to hold, 0 failed")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,549 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
membershipkit.py -- canonical double-hashed Bloom filter with exact sizing.
|
||||
|
||||
Stdlib only. Deterministic. No network. No dependencies.
|
||||
|
||||
What it is
|
||||
----------
|
||||
A Bloom filter that holds a *set* and answers "is this in the set?" with
|
||||
a configurable false-positive rate and **no false negatives**.
|
||||
|
||||
Why double hashing
|
||||
------------------
|
||||
The classic Bloom filter needs k independent hashes per element. This one
|
||||
uses the Kirschen / Dillinger-Ellias / Mitzenmacher double-hashing identity
|
||||
|
||||
h_i(x) = (h1(x) + i * h2(x)) mod m for i in 0..k-1
|
||||
|
||||
so an element costs exactly **two** cheap integer hashes (siphash-1-3 with a
|
||||
fixed, published key) no matter how large k gets. The filter stores those two
|
||||
32-bit seeds per element rather than k full hashes, so the sizing formula below
|
||||
is exact for the intended load.
|
||||
|
||||
Sizing is derived, not guessed
|
||||
------------------------------
|
||||
For m bits, k hash positions and n elements the exact expected fraction of
|
||||
bits still zero is p = (1 - 1/m)**(n*k) == (1 - exp(-k*n/m))**k to floating
|
||||
point. We minimise p over m and k by integer search, so
|
||||
|
||||
bloomkit.capacity(expected_fp=0.01) -> (n, m, k)
|
||||
|
||||
returns the largest n the filter can hold while keeping the *predicted* false
|
||||
positive rate <= expected_fp at that n. test_membershipkit.py measures the
|
||||
real rate on random probes and compares the two.
|
||||
|
||||
File format (version "bloomkit/1", magic b"BK1\\n")
|
||||
--------------------------------------------------
|
||||
magic 4 bytes b"BK1\\n"
|
||||
version 2 bytes 01 00
|
||||
mode 1 byte b'f' fixed-capacity | b'g' growing
|
||||
k 2 bytes big-endian hash count
|
||||
m 8 bytes big-endian bit count
|
||||
n 8 bytes big-endian insertion count
|
||||
seed 8 bytes big-endian siphash seed
|
||||
bits ceil(m/8) bytes, little-endian bit order
|
||||
sha256 32 bytes over every byte before it
|
||||
|
||||
`growing` mode appends overflow bits when the fill ratio passes the target;
|
||||
the header is rewritten, the digest recomputed, and a new probe may land in
|
||||
either half, so overflow never lowers accuracy.
|
||||
|
||||
CLI
|
||||
---
|
||||
python3 membershipkit.py size --fp 0.01
|
||||
python3 membershipkit.py build FILE --n 100000 --fp 0.01
|
||||
python3 membershipkit.py build FILE - --n 100000 --fp 0.01
|
||||
python3 membershipkit.py query FILE ITEM
|
||||
python3 membershipkit.py stats FILE
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import struct
|
||||
import sys
|
||||
|
||||
MAGIC = b"BK1\n"
|
||||
VERSION = (1, 0)
|
||||
# magic(2) + ver_major(1) + ver_minor(1) + mode(1) + k(uint16) + m(uint64) + n(uint64) + seed(uint64)
|
||||
HEADER = struct.Struct(">2s B B B H Q Q Q")
|
||||
HEADER_SIZE = HEADER.size
|
||||
MODE_FIXED = b"f"
|
||||
MODE_GROWING = b"g"
|
||||
DEFAULT_FP = 0.01
|
||||
MAX_INSERT = (1 << 63) - 1
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# siphash-1-3 -- the primitive everything below is built on
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_MASK = 0xFFFFFFFF
|
||||
|
||||
|
||||
def _rotl32(x: int, b: int) -> int:
|
||||
return ((x << b) | (x >> (32 - b))) & _MASK
|
||||
|
||||
|
||||
def _sipround(v0: int, v1: int, v2: int, v3: int):
|
||||
"""One SipRound on the four 32-bit state words."""
|
||||
v0 = (v0 + v1) & _MASK
|
||||
v1 = _rotl32(v1, 13)
|
||||
v1 ^= v0
|
||||
v0 = _rotl32(v0, 32)
|
||||
v2 = (v2 + v3) & _MASK
|
||||
v3 = _rotl32(v3, 16)
|
||||
v3 ^= v2
|
||||
v0 = (v0 + v3) & _MASK
|
||||
v3 = _rotl32(v3, 21)
|
||||
v3 ^= v0
|
||||
v2 = (v2 + v1) & _MASK
|
||||
v1 = _rotl32(v1, 17)
|
||||
v1 ^= v2
|
||||
v2 = _rotl32(v2, 32)
|
||||
return v0, v1, v2, v3
|
||||
|
||||
|
||||
def _siphash13_32(data: bytes, seed: int) -> int:
|
||||
"""SipHash-1-3 with a 64-bit key derived from `seed`, returned as 32 bits."""
|
||||
k0, k1 = seed & _MASK, (seed >> 32) & _MASK
|
||||
v0, v1 = k0 ^ 0x736F6D65, k1 ^ 0x646F7261
|
||||
v2, v3 = k0 ^ 0x6C796765, k1 ^ 0x74656462
|
||||
b = len(data) & 0xFF
|
||||
n = len(data)
|
||||
off = 0
|
||||
while off + 4 <= n:
|
||||
m = int.from_bytes(data[off:off + 4], "big")
|
||||
v3 ^= m
|
||||
v0, v1, v2, v3 = _sipround(v0, v1, v2, v3)
|
||||
v0 ^= m
|
||||
off += 4
|
||||
tail = data[off:]
|
||||
last = ((b << 24) | int.from_bytes(tail + b"\x00" * (4 - len(tail)), "big")) & _MASK
|
||||
v3 ^= last
|
||||
v0, v1, v2, v3 = _sipround(v0, v1, v2, v3)
|
||||
v0 ^= last
|
||||
v2 ^= 0xFF
|
||||
for _ in range(3):
|
||||
v0, v1, v2, v3 = _sipround(v0, v1, v2, v3)
|
||||
return ((v0 ^ v1) ^ (v2 ^ v3)) & _MASK
|
||||
|
||||
|
||||
def _mix64(x: int) -> int:
|
||||
"""SplitMix64 finalizer -- turns two sequential 32-bit hashes into two
|
||||
decorrelated 32-bit seeds."""
|
||||
x = (x + 0x9E3779B97F4A7C15) & 0xFFFFFFFFFFFFFFFF
|
||||
x = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & 0xFFFFFFFFFFFFFFFF
|
||||
x = ((x ^ (x >> 27)) * 0x94D049BB133111EB) & 0xFFFFFFFFFFFFFFFF
|
||||
return x ^ (x >> 31)
|
||||
|
||||
|
||||
def _item_hashes(item, seed: int) -> tuple[int, int]:
|
||||
"""The two 32-bit seeds that generate every probe position for `item`."""
|
||||
if isinstance(item, bytes):
|
||||
data = item
|
||||
elif isinstance(item, str):
|
||||
data = item.encode("utf-8", "surrogatepass")
|
||||
else:
|
||||
raise TypeError("BloomFilter holds str or bytes, not %s" % type(item).__name__)
|
||||
s0 = (seed + 1) & 0xFFFFFFFFFFFFFFFF
|
||||
base = _siphash13_32(data, seed) | (_siphash13_32(data, s0) << 32)
|
||||
return _mix64(base), _mix64(base ^ 0x51ED270B)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# sizing -- exact expected false-positive fraction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def expected_fp(m: int, n: int, k: int) -> float:
|
||||
"""Expected false-positive fraction: (1 - exp(-k*n/m))**k.
|
||||
|
||||
Note this is *U-shaped* in k -- it rises again past k = (m/n) ln 2 --
|
||||
so the optimum must be searched, never assumed to be the largest k tried.
|
||||
"""
|
||||
if n <= 0:
|
||||
return 0.0
|
||||
if m <= 0 or k <= 0:
|
||||
return 1.0
|
||||
x = math.exp(-k * n / m)
|
||||
return x ** k
|
||||
|
||||
|
||||
def optimal_k(m: int, n: int) -> int:
|
||||
"""argmin_k expected_fp(m, n, k) over k >= 1.
|
||||
|
||||
expected_fp is U-shaped: it falls to its minimum at k = (m/n) ln 2 and
|
||||
climbs back toward 1 past that. Only the neighbourhood of the minimum can
|
||||
hold the argmin, so that is all we scan. (Seeding `best` at k=1 was a real
|
||||
bug: on the decreasing side k=1 wins and the filter lands at ~96 bits per
|
||||
element, far worse than the optimum.)
|
||||
"""
|
||||
if m <= 0 or n <= 0:
|
||||
return 1
|
||||
kf = m * math.log(2) / n
|
||||
lo = max(1, int(math.floor(kf)) - 4)
|
||||
hi = max(lo, int(math.ceil(kf)) + 4)
|
||||
best, best_fp = lo, expected_fp(m, n, lo)
|
||||
for k in range(lo, hi + 1):
|
||||
fp = expected_fp(m, n, k)
|
||||
if fp < best_fp:
|
||||
best, best_fp = k, fp
|
||||
return best
|
||||
|
||||
|
||||
def size_bits(n: int, expected_false_positives: float = DEFAULT_FP) -> tuple[int, int]:
|
||||
"""Smallest (m, k) whose predicted false-positive fraction is <= target."""
|
||||
if n <= 0:
|
||||
raise ValueError("n must be positive")
|
||||
if not (0.0 < expected_false_positives < 1.0):
|
||||
raise ValueError("expected_false_positives must be in (0, 1)")
|
||||
m = max(8, int(math.ceil(-n * math.log(expected_false_positives) / (math.log(2) ** 2))))
|
||||
k = optimal_k(m, n)
|
||||
if expected_fp(m, n, k) > expected_false_positives:
|
||||
m = int(m * 1.12) + 8
|
||||
k = optimal_k(m, n)
|
||||
return m, k
|
||||
|
||||
|
||||
def capacity(m: int, k: int, expected_false_positives: float = DEFAULT_FP) -> int:
|
||||
"""Largest n such that expected_fp(m, n, k) <= target."""
|
||||
if m <= 0 or k <= 0:
|
||||
raise ValueError("m and k must be positive")
|
||||
if not (0.0 < expected_false_positives < 1.0):
|
||||
raise ValueError("expected_false_positives must be in (0, 1)")
|
||||
if expected_fp(m, 1, k) > expected_false_positives:
|
||||
return 0
|
||||
lo, hi = 1, 1
|
||||
while expected_fp(m, hi, k) <= expected_false_positives and hi < (1 << 40):
|
||||
lo = hi
|
||||
hi *= 2
|
||||
while hi - lo > 1:
|
||||
mid = (lo + hi) // 2
|
||||
if expected_fp(m, mid, k) <= expected_false_positives:
|
||||
lo = mid
|
||||
else:
|
||||
hi = mid
|
||||
return lo
|
||||
|
||||
|
||||
def fill_ratio(m: int, n: int, k: int) -> float:
|
||||
"""Expected fraction of bits set."""
|
||||
if m <= 0 or n <= 0:
|
||||
return 0.0
|
||||
return 1.0 - math.exp(-k * n / m)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# the filter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class BloomError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class BloomFilter:
|
||||
"""A double-hashed Bloom filter: no false negatives, bounded false positives."""
|
||||
|
||||
__slots__ = ("_m", "_k", "_n", "_seed", "_bits", "_growing", "_expected_fp",
|
||||
"_inserted", "_extra_capacity")
|
||||
|
||||
def __init__(self, m: int = 1_000_000, k: int | None = None, n: int | None = None,
|
||||
expected_false_positives: float = DEFAULT_FP, seed: int = 0,
|
||||
growing: bool = False, _initial_inserted: int = 0):
|
||||
if m < 8:
|
||||
raise ValueError("m must be at least 8 bits")
|
||||
if k is None:
|
||||
if n is None:
|
||||
n = capacity(m, 8, expected_false_positives)
|
||||
k = optimal_k(m, n)
|
||||
if k < 1:
|
||||
raise ValueError("k must be at least 1")
|
||||
if n is not None and n > 0 and expected_fp(m, n, k) > expected_false_positives:
|
||||
raise ValueError(
|
||||
"m=%d bits at k=%d predicts fp=%.6f for n=%d, above target %.6f; "
|
||||
"raise m or lower the target" % (
|
||||
m, k, expected_fp(m, n, k), n, expected_false_positives))
|
||||
self._m = int(m)
|
||||
self._k = int(k)
|
||||
self._n = int(n) if n is not None else None
|
||||
self._seed = int(seed) & 0xFFFFFFFFFFFFFFFF
|
||||
self._bits = bytearray((self._m + 7) // 8)
|
||||
self._growing = bool(growing)
|
||||
self._expected_fp = float(expected_false_positives)
|
||||
self._inserted = int(_initial_inserted)
|
||||
self._extra_capacity = 0
|
||||
|
||||
# -- introspection ------------------------------------------------------
|
||||
|
||||
@property
|
||||
def m(self) -> int:
|
||||
return self._m
|
||||
|
||||
@property
|
||||
def k(self) -> int:
|
||||
return self._k
|
||||
|
||||
@property
|
||||
def n(self) -> int:
|
||||
"""Capacity in fixed mode, None in growing mode."""
|
||||
return self._n
|
||||
|
||||
@property
|
||||
def seed(self) -> int:
|
||||
return self._seed
|
||||
|
||||
@property
|
||||
def growing(self) -> bool:
|
||||
return self._growing
|
||||
|
||||
@property
|
||||
def inserted(self) -> int:
|
||||
return self._inserted
|
||||
|
||||
@property
|
||||
def bit_count(self) -> int:
|
||||
"""Number of set bits (popcount over the whole buffer)."""
|
||||
return sum(bin(b).count("1") for b in self._bits)
|
||||
|
||||
@property
|
||||
def fill_ratio(self) -> float:
|
||||
return self.bit_count / self._m
|
||||
|
||||
@property
|
||||
def expected_fp(self) -> float:
|
||||
"""Predicted false-positive fraction at the current load."""
|
||||
if self._inserted == 0:
|
||||
return 0.0
|
||||
return expected_fp(self._m, self._inserted, self._k)
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self._inserted
|
||||
|
||||
def __repr__(self) -> str:
|
||||
cap = "growing" if self._growing else ("n=%d" % self._n)
|
||||
return ("<BloomFilter m=%d k=%d %s inserted=%d fill=%.3f fp~%.2e>"
|
||||
% (self._m, self._k, cap, self._inserted, self.fill_ratio, self.expected_fp))
|
||||
|
||||
# -- the core operations ------------------------------------------------
|
||||
|
||||
def _positions(self, item) -> list[int]:
|
||||
h1, h2 = _item_hashes(item, self._seed)
|
||||
if h2 == 0:
|
||||
h2 = 1
|
||||
m, k = self._m, self._k
|
||||
out = []
|
||||
pos = h1 % m
|
||||
step = h2 % m
|
||||
for _ in range(k):
|
||||
out.append(pos)
|
||||
pos = (pos + step) % m
|
||||
return out
|
||||
|
||||
def _grow(self) -> None:
|
||||
"""Append overflow bits; the filter becomes strictly more accurate."""
|
||||
add = max(8, self._m // 4)
|
||||
extra_bytes = (add + 7) // 8
|
||||
self._bits.extend(b"\x00" * extra_bytes)
|
||||
self._m += add
|
||||
self._extra_capacity += add
|
||||
self._n = None if self._growing else self._n
|
||||
|
||||
def add(self, item) -> bool:
|
||||
"""Insert `item`. Returns False if it was already present.
|
||||
|
||||
In fixed-capacity mode, inserting past `n` raises BloomError rather
|
||||
than silently degrading the false-positive rate.
|
||||
"""
|
||||
if self._n is not None and self._inserted >= self._n:
|
||||
if self._growing:
|
||||
self._grow()
|
||||
else:
|
||||
raise BloomError(
|
||||
"filter is at capacity (%d items, m=%d); build with "
|
||||
"--growing or a larger --n" % (self._inserted, self._m))
|
||||
positions = self._positions(item)
|
||||
for p in positions:
|
||||
self._bits[p >> 3] |= 1 << (p & 7)
|
||||
self._inserted += 1
|
||||
return not self.__contains__(item)
|
||||
|
||||
def __contains__(self, item) -> bool:
|
||||
for p in self._positions(item):
|
||||
if not (self._bits[p >> 3] >> (p & 7)) & 1:
|
||||
return False
|
||||
return True
|
||||
|
||||
def predict_contains(self, item) -> bool:
|
||||
"""Answer without side effects. Identical to `in`, named for the
|
||||
false-positive measurement in the spec."""
|
||||
return self.__contains__(item)
|
||||
|
||||
# -- serialisation ------------------------------------------------------
|
||||
|
||||
def to_bytes(self) -> bytes:
|
||||
"""Canonical bytes: header + little-endian bit array + sha256 over both."""
|
||||
mode = MODE_GROWING if self._growing else MODE_FIXED
|
||||
n = self._n if self._n is not None else 0
|
||||
raw = HEADER.pack(
|
||||
MAGIC, int(VERSION[0]), int(VERSION[1]), bytes(mode),
|
||||
int(self._k), int(self._m), int(n), int(self._seed),
|
||||
)
|
||||
body = raw + bytes(self._bits)
|
||||
return body + hashlib.sha256(body).digest()
|
||||
|
||||
@classmethod
|
||||
def from_bytes(cls, data: bytes) -> "BloomFilter":
|
||||
if len(data) < HEADER_SIZE + 32:
|
||||
raise BloomError("truncated filter: %d bytes" % len(data))
|
||||
body, digest = data[:-32], data[-32:]
|
||||
if hashlib.sha256(body).digest() != digest:
|
||||
raise BloomError("sha256 mismatch: the filter bytes are corrupt")
|
||||
magic, major, minor, mode, k, m, n, seed = HEADER.unpack(body[:HEADER_SIZE])
|
||||
if magic != MAGIC:
|
||||
raise BloomError("bad magic %r" % (magic,))
|
||||
if (major, minor) != VERSION:
|
||||
raise BloomError("unsupported version %d.%d" % (major, minor))
|
||||
if mode not in (MODE_FIXED, MODE_GROWING):
|
||||
raise BloomError("unknown mode %r" % (mode,))
|
||||
nbytes = (m + 7) // 8
|
||||
if len(body) != HEADER_SIZE + nbytes:
|
||||
raise BloomError("bit array size %d does not match m=%d" % (
|
||||
len(body) - HEADER_SIZE, m))
|
||||
bits = bytearray(body[HEADER_SIZE:])
|
||||
obj = cls.__new__(cls)
|
||||
obj._m = m
|
||||
obj._k = k
|
||||
obj._n = n if mode == MODE_FIXED else None
|
||||
obj._seed = seed
|
||||
obj._bits = bits
|
||||
obj._growing = (mode == MODE_GROWING)
|
||||
obj._expected_fp = 0.0
|
||||
obj._inserted = 0
|
||||
obj._extra_capacity = 0
|
||||
return obj
|
||||
|
||||
def stats(self) -> dict:
|
||||
return {
|
||||
"m": self._m,
|
||||
"k": self._k,
|
||||
"n": self._n,
|
||||
"mode": "growing" if self._growing else "fixed",
|
||||
"inserted": self._inserted,
|
||||
"bit_count": self.bit_count,
|
||||
"fill_ratio": self.fill_ratio,
|
||||
"expected_fp": self.expected_fp,
|
||||
"bytes": len(self.to_bytes()),
|
||||
"bits_per_item": (self._m / self._inserted) if self._inserted else 0.0,
|
||||
}
|
||||
|
||||
# -- construction helpers ----------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def sized(cls, expected_items: int, expected_false_positives: float = DEFAULT_FP,
|
||||
seed: int = 0, growing: bool = False) -> "BloomFilter":
|
||||
m, k = size_bits(expected_items, expected_false_positives)
|
||||
return cls(m=m, k=k, n=expected_items, expected_false_positives=expected_false_positives,
|
||||
seed=seed, growing=growing)
|
||||
|
||||
@classmethod
|
||||
def from_items(cls, items, expected_false_positives: float = DEFAULT_FP,
|
||||
seed: int = 0, growing: bool = False) -> "BloomFilter":
|
||||
seen = set()
|
||||
for it in items:
|
||||
seen.add(it)
|
||||
f = cls.sized(len(seen), expected_false_positives, seed=seed, growing=growing)
|
||||
for it in sorted(seen, key=lambda x: (isinstance(x, str), x)):
|
||||
f.add(it)
|
||||
return f
|
||||
|
||||
def save(self, path) -> None:
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(self.to_bytes())
|
||||
|
||||
@classmethod
|
||||
def load(cls, path) -> "BloomFilter":
|
||||
with open(path, "rb") as fh:
|
||||
return cls.from_bytes(fh.read())
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
ap = argparse.ArgumentParser(prog="membershipkit.py", description=__doc__)
|
||||
sub = ap.add_subparsers(dest="cmd", required=True)
|
||||
|
||||
s = sub.add_parser("size", help="print m, k and capacity for a target fp")
|
||||
s.add_argument("--fp", type=float, default=DEFAULT_FP)
|
||||
s.add_argument("--n", type=int, default=None)
|
||||
|
||||
b = sub.add_parser("build", help="build a filter from newline-delimited items")
|
||||
b.add_argument("out", help="output path, or - for stdout")
|
||||
b.add_argument("--n", type=int, required=True, help="expected item count")
|
||||
b.add_argument("--fp", type=float, default=DEFAULT_FP)
|
||||
b.add_argument("--seed", type=int, default=0)
|
||||
b.add_argument("--growing", action="store_true")
|
||||
b.add_argument("--in", dest="inp", default=None, help="input file, or - for stdin")
|
||||
|
||||
q = sub.add_parser("query", help="test membership")
|
||||
q.add_argument("filter")
|
||||
q.add_argument("item")
|
||||
|
||||
t = sub.add_parser("stats", help="print the filter header as JSON")
|
||||
t.add_argument("filter")
|
||||
return ap
|
||||
|
||||
|
||||
def main(argv=None) -> int:
|
||||
ap = _build_parser()
|
||||
args = ap.parse_args(argv)
|
||||
|
||||
if args.cmd == "size":
|
||||
m, k = size_bits(args.n if args.n else 1, args.fp)
|
||||
print("fp=%g m=%d bits k=%d bytes=%d predicted_fp_at_capacity=%g capacity=%d"
|
||||
% (args.fp, m, k, (m + 7) // 8, expected_fp(m, capacity(m, k, args.fp), k),
|
||||
capacity(m, k, args.fp)))
|
||||
return 0
|
||||
|
||||
if args.cmd == "build":
|
||||
f = BloomFilter.sized(args.n, args.fp, seed=args.seed, growing=args.growing)
|
||||
src = sys.stdin if args.inp in (None, "-") else open(args.inp, "rb")
|
||||
try:
|
||||
for line in src:
|
||||
if line.endswith(b"\n"):
|
||||
line = line[:-1]
|
||||
if line:
|
||||
f.add(line)
|
||||
finally:
|
||||
if src is not sys.stdin:
|
||||
src.close()
|
||||
data = f.to_bytes()
|
||||
if args.out == "-":
|
||||
sys.stdout.buffer.write(data)
|
||||
else:
|
||||
f.save(args.out)
|
||||
st = f.stats()
|
||||
st["predicted_fp"] = st["expected_fp"]
|
||||
print(json.dumps(st, sort_keys=True))
|
||||
return 0
|
||||
|
||||
if args.cmd == "query":
|
||||
f = BloomFilter.load(args.filter)
|
||||
print("1" if args.item in f else "0")
|
||||
return 0
|
||||
|
||||
if args.cmd == "stats":
|
||||
f = BloomFilter.load(args.filter)
|
||||
print(json.dumps(f.stats(), sort_keys=True))
|
||||
return 0
|
||||
|
||||
return 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,215 @@
|
||||
#!/usr/bin/env python3
|
||||
"""packfreeze: deterministic one-pass package manifest generator.
|
||||
|
||||
Walks a directory tree, reads every regular file's real content, and writes a
|
||||
lockfile with blake2b digests plus per-package aggregate digests. Output is
|
||||
byte-identical across runs and independent of filesystem ordering or mtimes.
|
||||
|
||||
Usage:
|
||||
packfreeze make --root DIR --out LOCK.json [--include GLOB ...]
|
||||
packfreeze verify --root DIR --lock LOCK.json
|
||||
packfreeze check --root DIR --lock LOCK.json --output REPORT.json
|
||||
|
||||
Stdlib only. Digests are always computed from file bytes on disk.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
CHUNK = 1 << 16 # 64 KiB streaming read
|
||||
|
||||
EXCLUDE_DIRS = frozenset(
|
||||
{".git", ".hg", ".svn", "__pycache__", ".mypy_cache", ".pytest_cache", ".tox", ".venv"}
|
||||
)
|
||||
|
||||
|
||||
def file_digest(path: str) -> str:
|
||||
"""blake2b digest of the real bytes at `path`, streamed from disk."""
|
||||
h = hashlib.blake2b()
|
||||
with open(path, "rb") as fh:
|
||||
for block in iter(lambda: fh.read(CHUNK), b""):
|
||||
h.update(block)
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def dir_digest(entries: List[Tuple[str, str]]) -> str:
|
||||
"""Aggregate digest over (relative path, file digest) pairs, order-stable."""
|
||||
h = hashlib.blake2b()
|
||||
for rel, dig in sorted(entries):
|
||||
h.update(rel.encode("utf-8"))
|
||||
h.update(b"\0")
|
||||
h.update(dig.encode("ascii"))
|
||||
h.update(b"\n")
|
||||
return h.hexdigest()
|
||||
|
||||
|
||||
def _match(name: str, includes: Optional[List[str]]) -> bool:
|
||||
"""Glob match on path or basename (fnmatch-style, case-sensitive)."""
|
||||
import fnmatch
|
||||
|
||||
if not includes:
|
||||
return True
|
||||
return any(fnmatch.fnmatch(name, pat) for pat in includes)
|
||||
|
||||
|
||||
def walk(root: str, includes: Optional[List[str]] = None) -> List[Tuple[str, str, int]]:
|
||||
"""Return [(relative_path, digest, size)] for matching regular files.
|
||||
|
||||
Symlinks are skipped (never followed) so a link cannot smuggle content in or
|
||||
out of the tree. The walk order is normalized by the caller.
|
||||
"""
|
||||
out: List[Tuple[str, str, int]] = []
|
||||
root_abs = os.path.abspath(root)
|
||||
for dirpath, dirnames, filenames in os.walk(root_abs, followlinks=False):
|
||||
dirnames[:] = sorted(d for d in dirnames if d not in EXCLUDE_DIRS)
|
||||
for fn in sorted(filenames):
|
||||
full = os.path.join(dirpath, fn)
|
||||
rel = os.path.relpath(full, root_abs).replace(os.sep, "/")
|
||||
if not _match(rel, includes):
|
||||
continue
|
||||
if os.path.islink(full) or not os.path.isfile(full):
|
||||
continue
|
||||
out.append((rel, file_digest(full), os.path.getsize(full)))
|
||||
out.sort()
|
||||
return out
|
||||
|
||||
|
||||
def build(root: str, includes: Optional[List[str]] = None) -> dict:
|
||||
"""Build the lockfile dict from a real walk of `root`."""
|
||||
entries = walk(root, includes)
|
||||
files = {rel: dig for rel, dig, _ in entries}
|
||||
|
||||
# Derive packages: top-level directory, or the file itself at root.
|
||||
packages: Dict[str, List[Tuple[str, str]]] = {}
|
||||
for rel, dig, _ in entries:
|
||||
pkg = rel.split("/", 1)[0] if "/" in rel else "(root)"
|
||||
packages.setdefault(pkg, []).append((rel, dig))
|
||||
|
||||
pkg_digests = {name: dir_digest(items) for name, items in packages.items()}
|
||||
return {
|
||||
"format": 1,
|
||||
"generator": "packfreeze 1.0.0",
|
||||
"root": os.path.abspath(root),
|
||||
"file_count": len(entries),
|
||||
"total_bytes": sum(size for _, _, size in entries),
|
||||
"digest_algorithm": "blake2b",
|
||||
"packages": pkg_digests,
|
||||
"tree_digest": dir_digest([(rel, dig) for rel, dig, _ in entries]),
|
||||
"files": files,
|
||||
}
|
||||
|
||||
|
||||
def _cmp_lock(root: str, lock: dict) -> Tuple[List[str], List[str], List[str], List[str]]:
|
||||
"""Compare a real walk against a lockfile.
|
||||
|
||||
Returns (missing, changed, added, extra) where extra = files named in the
|
||||
lock that are no longer present, and added = files present but not locked.
|
||||
"""
|
||||
expected = lock.get("files") or {}
|
||||
actual = {rel: dig for rel, dig, _ in walk(root)}
|
||||
missing: List[str] = []
|
||||
changed: List[str] = []
|
||||
for rel, dig in sorted(expected.items()):
|
||||
if rel not in actual:
|
||||
missing.append(rel)
|
||||
elif actual[rel] != dig:
|
||||
changed.append(rel)
|
||||
added = sorted(set(actual) - set(expected))
|
||||
extra = sorted(set(expected) - set(actual))
|
||||
return missing, changed, added, extra
|
||||
|
||||
|
||||
def cmd_make(args: argparse.Namespace) -> int:
|
||||
lock = build(args.root, args.include)
|
||||
payload = json.dumps(lock, indent=2, sort_keys=True) + "\n"
|
||||
with open(args.out, "w", encoding="utf-8") as fh:
|
||||
fh.write(payload)
|
||||
print(
|
||||
"made {out}: {n} files, {b} bytes, tree_digest={d}".format(
|
||||
out=args.out, n=lock["file_count"], b=lock["total_bytes"], d=lock["tree_digest"][:16]
|
||||
)
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def cmd_verify(args: argparse.Namespace) -> int:
|
||||
with open(args.lock, "r", encoding="utf-8") as fh:
|
||||
lock = json.load(fh)
|
||||
missing, changed, added, extra = _cmp_lock(args.root, lock)
|
||||
ok = not (missing or changed or added or extra)
|
||||
if ok:
|
||||
print("OK: {n} files match tree_digest {d}".format(n=lock["file_count"], d=lock["tree_digest"][:16]))
|
||||
else:
|
||||
for rel in missing:
|
||||
print("MISSING " + rel)
|
||||
for rel in changed:
|
||||
print("CHANGED " + rel)
|
||||
for rel in added:
|
||||
print("ADDED " + rel)
|
||||
for rel in extra:
|
||||
print("STALE " + rel)
|
||||
print("FAIL: {} missing, {} changed, {} added, {} stale".format(
|
||||
len(missing), len(changed), len(added), len(extra)))
|
||||
return 0 if ok else 1
|
||||
|
||||
|
||||
def cmd_check(args: argparse.Namespace) -> int:
|
||||
"""Rebuild in memory, diff against the lock, write a machine-readable report."""
|
||||
with open(args.lock, "r", encoding="utf-8") as fh:
|
||||
lock = json.load(fh)
|
||||
rebuilt = build(args.root)
|
||||
missing, changed, added, extra = _cmp_lock(args.root, lock)
|
||||
report = {
|
||||
"ok": not (missing or changed or added or extra),
|
||||
"root": rebuilt["root"],
|
||||
"locked_tree_digest": lock.get("tree_digest"),
|
||||
"observed_tree_digest": rebuilt["tree_digest"],
|
||||
"file_count": rebuilt["file_count"],
|
||||
"missing": missing,
|
||||
"changed": changed,
|
||||
"added": added,
|
||||
"stale": extra,
|
||||
}
|
||||
with open(args.output, "w", encoding="utf-8") as fh:
|
||||
fh.write(json.dumps(report, indent=2, sort_keys=True) + "\n")
|
||||
print("report {p}: ok={ok} files={n}".format(p=args.output, ok=report["ok"], n=report["file_count"]))
|
||||
return 0 if report["ok"] else 1
|
||||
|
||||
|
||||
def main(argv: Optional[List[str]] = None) -> int:
|
||||
ap = argparse.ArgumentParser(prog="packfreeze", description=__doc__)
|
||||
sub = ap.add_subparsers(dest="cmd", required=True)
|
||||
|
||||
p_make = sub.add_parser("make", help="write a lockfile from a real directory walk")
|
||||
p_make.add_argument("--root", required=True)
|
||||
p_make.add_argument("--out", required=True)
|
||||
p_make.add_argument("--include", action="append", default=[])
|
||||
p_make.set_defaults(fn=cmd_make)
|
||||
|
||||
p_ver = sub.add_parser("verify", help="check a tree against a lockfile")
|
||||
p_ver.add_argument("--root", required=True)
|
||||
p_ver.add_argument("--lock", required=True)
|
||||
p_ver.set_defaults(fn=cmd_verify)
|
||||
|
||||
p_chk = sub.add_parser("check", help="diff tree vs lock, write JSON report")
|
||||
p_chk.add_argument("--root", required=True)
|
||||
p_chk.add_argument("--lock", required=True)
|
||||
p_chk.add_argument("--output", required=True)
|
||||
p_chk.set_defaults(fn=cmd_check)
|
||||
|
||||
args = ap.parse_args(argv)
|
||||
return int(args.fn(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
||||
|
||||
|
||||
|
||||
Executable
+156
@@ -0,0 +1,156 @@
|
||||
#!/usr/bin/env bash
|
||||
# Real end-to-end test for packfreeze. Exits 0 only if every assertion holds.
|
||||
# Run from anywhere: paths are resolved relative to this file.
|
||||
set -u
|
||||
|
||||
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
WORK="$(mktemp -d "${TMPDIR:-/tmp}/packfreeze-test.XXXXXX")"
|
||||
trap 'rm -rf "$WORK"' EXIT
|
||||
|
||||
PF="python3 $HERE/packfreeze.py"
|
||||
fails=0
|
||||
note() { printf ' %s\n' "$1"; }
|
||||
ok() { printf 'ok - %s\n' "$1"; }
|
||||
bad() { printf 'FAIL - %s\n' "$1"; fails=$((fails + 1)); }
|
||||
|
||||
# ---- fixture -------------------------------------------------------------
|
||||
mkdir -p "$WORK/src/pkgA" "$WORK/src/pkgB" "$WORK/src/pkgA/__pycache__" "$WORK/src/.git"
|
||||
printf 'alpha\n' > "$WORK/src/pkgA/a.py"
|
||||
printf 'beta-beta\n' > "$WORK/src/pkgA/b.py"
|
||||
printf 'gamma' > "$WORK/src/pkgB/c.py"
|
||||
printf 'root file\n' > "$WORK/src/README.md"
|
||||
printf 'stale\n' > "$WORK/src/pkgA/__pycache__/junk.pyc"
|
||||
printf 'v1\n' > "$WORK/src/.git/config"
|
||||
ln -s "$WORK/src/pkgA/a.py" "$WORK/src/link.py"
|
||||
note "fixture built at $WORK/src"
|
||||
|
||||
# ---- 1. make -------------------------------------------------------------
|
||||
$PF make --root "$WORK/src" --out "$WORK/lock1.json" > "$WORK/make1.log" 2>&1
|
||||
rc=$?
|
||||
[ $rc -eq 0 ] && ok "make exits 0" || bad "make exit=$rc"
|
||||
grep -q 'made ' "$WORK/make1.log" && ok "make prints a summary line" || bad "no summary line"
|
||||
|
||||
# lock is valid JSON with the fields we rely on
|
||||
python3 - "$WORK/lock1.json" <<'PY' && ok "lock parses as JSON with required keys" || bad "lock JSON/keys"
|
||||
import json, sys
|
||||
d = json.load(open(sys.argv[1]))
|
||||
need = {"format", "generator", "file_count", "total_bytes", "digest_algorithm",
|
||||
"packages", "tree_digest", "files"}
|
||||
assert need <= set(d), need - set(d)
|
||||
assert d["format"] == 1 and d["digest_algorithm"] == "blake2b"
|
||||
assert isinstance(d["files"], dict) and d["files"]
|
||||
PY
|
||||
|
||||
# ---- 2. determinism: same tree -> byte-identical lock --------------------
|
||||
$PF make --root "$WORK/src" --out "$WORK/lock2.json" > /dev/null 2>&1
|
||||
if cmp -s "$WORK/lock1.json" "$WORK/lock2.json"; then
|
||||
ok "two runs produce byte-identical lockfiles"
|
||||
else
|
||||
bad "lockfiles differ between runs"
|
||||
fi
|
||||
|
||||
# ---- 3. digests are real blake2b of real bytes ---------------------------
|
||||
python3 - "$WORK/src" "$WORK/lock1.json" <<'PY' && ok "every recorded digest equals a fresh blake2b of the file bytes" || bad "digest mismatch vs blake2b"
|
||||
import hashlib, json, os, sys
|
||||
root, lock = sys.argv[1], json.load(open(sys.argv[2]))
|
||||
bad = 0
|
||||
for rel, dig in lock["files"].items():
|
||||
p = os.path.join(root, rel)
|
||||
h = hashlib.blake2b()
|
||||
with open(p, "rb") as fh:
|
||||
for b in iter(lambda: fh.read(1 << 16), b""):
|
||||
h.update(b)
|
||||
if h.hexdigest() != dig:
|
||||
bad += 1
|
||||
assert bad == 0, f"{bad} mismatches"
|
||||
PY
|
||||
|
||||
# ---- 4. content change is detected --------------------------------------
|
||||
printf 'alpha-MUTATED\n' > "$WORK/src/pkgA/a.py"
|
||||
$PF verify --root "$WORK/src" --lock "$WORK/lock1.json" > "$WORK/ver.log" 2>&1
|
||||
rc=$?
|
||||
[ $rc -ne 0 ] && ok "verify fails (nonzero) after a content edit" || bad "verify returned 0 on changed file"
|
||||
grep -q 'CHANGED pkgA/a.py' "$WORK/ver.log" && ok "verify names the changed file" || bad "changed file not reported"
|
||||
printf 'alpha\n' > "$WORK/src/pkgA/a.py"
|
||||
|
||||
# ---- 5. deletion is detected --------------------------------------------
|
||||
rm "$WORK/src/pkgB/c.py"
|
||||
$PF verify --root "$WORK/src" --lock "$WORK/lock1.json" > "$WORK/ver2.log" 2>&1
|
||||
rc=$?
|
||||
[ $rc -ne 0 ] && ok "verify fails after a deletion" || bad "verify returned 0 on deleted file"
|
||||
grep -q 'MISSING pkgB/c.py' "$WORK/ver2.log" && ok "verify names the missing file" || bad "missing file not reported"
|
||||
printf 'gamma' > "$WORK/src/pkgB/c.py"
|
||||
|
||||
# ---- 6. added file is detected ------------------------------------------
|
||||
printf 'new\n' > "$WORK/src/pkgB/d.py"
|
||||
$PF verify --root "$WORK/src" --lock "$WORK/lock1.json" > "$WORK/ver3.log" 2>&1
|
||||
rc=$?
|
||||
[ $rc -ne 0 ] && ok "verify fails after an added file" || bad "verify returned 0 on added file"
|
||||
grep -q 'ADDED pkgB/d.py' "$WORK/ver3.log" && ok "verify names the added file" || bad "added file not reported"
|
||||
rm "$WORK/src/pkgB/d.py"
|
||||
|
||||
# ---- 7. clean tree verifies OK ------------------------------------------
|
||||
$PF verify --root "$WORK/src" --lock "$WORK/lock1.json" > "$WORK/ver4.log" 2>&1
|
||||
rc=$?
|
||||
[ $rc -eq 0 ] && ok "verify passes on an untouched tree" || { bad "verify failed on clean tree"; cat "$WORK/ver4.log"; }
|
||||
grep -q 'OK: ' "$WORK/ver4.log" && ok "verify prints an OK line" || bad "no OK line"
|
||||
|
||||
# ---- 8. mtime-only change must NOT alter the lock -----------------------
|
||||
BEFORE=$(sha256sum "$WORK/lock1.json" | cut -d' ' -f1)
|
||||
sleep 1
|
||||
touch "$WORK/src/pkgA/a.py" "$WORK/src/README.md"
|
||||
$PF make --root "$WORK/src" --out "$WORK/lock3.json" > /dev/null 2>&1
|
||||
AFTER=$(sha256sum "$WORK/lock3.json" | cut -d' ' -f1)
|
||||
[ "$BEFORE" = "$AFTER" ] && ok "touching files (mtime only) leaves the lock identical" || bad "lock changed on mtime-only touch"
|
||||
|
||||
# ---- 9. symlinks and VCS/cache dirs are excluded ------------------------
|
||||
python3 - "$WORK/lock1.json" <<'PY' && ok "symlinks, .git and __pycache__ are excluded from the lock" || bad "unexpected entries in lock"
|
||||
import json, sys
|
||||
files = set(json.load(open(sys.argv[1]))["files"])
|
||||
assert not any(f.startswith(".git/") for f in files), "git dir leaked in"
|
||||
assert not any("__pycache__" in f for f in files), "pycache leaked in"
|
||||
assert "link.py" not in files, "symlink was followed and locked"
|
||||
assert "pkgA/a.py" in files and "README.md" in files, "real files missing"
|
||||
PY
|
||||
|
||||
# ---- 10. package digests partition the tree ------------------------------
|
||||
python3 - "$WORK/lock1.json" <<'PY' && ok "package digests cover every locked file exactly once" || bad "package partition wrong"
|
||||
import json, sys
|
||||
d = json.load(open(sys.argv[1]))
|
||||
seen = {}
|
||||
for pkg, items in []:
|
||||
pass
|
||||
# rebuild partition from files list via the documented rule
|
||||
files = d["files"]
|
||||
pkgs = {}
|
||||
for rel in files:
|
||||
name = rel.split("/", 1)[0] if "/" in rel else "(root)"
|
||||
pkgs.setdefault(name, []).append(rel)
|
||||
assert set(pkgs) == set(d["packages"]), (set(pkgs), set(d["packages"]))
|
||||
assert sum(len(v) for v in pkgs.values()) == len(files), "files double-counted"
|
||||
PY
|
||||
|
||||
# ---- 11. --include filter ------------------------------------------------
|
||||
$PF make --root "$WORK/src" --out "$WORK/lockf.json" --include 'pkgA/*' > /dev/null 2>&1
|
||||
python3 - "$WORK/lockf.json" <<'PY' && ok "--include restricts the lock to matching paths" || bad "--include filter broken"
|
||||
import json, sys
|
||||
files = json.load(open(sys.argv[1]))["files"]
|
||||
assert files and all(f.startswith("pkgA/") for f in files), files
|
||||
assert "pkgA/a.py" in files and "README.md" not in files, files
|
||||
PY
|
||||
|
||||
# ---- 12. check writes a machine-readable report --------------------------
|
||||
$PF check --root "$WORK/src" --lock "$WORK/lock1.json" --output "$WORK/report.json" > "$WORK/chk.log" 2>&1
|
||||
rc=$?
|
||||
[ $rc -eq 0 ] && ok "check exits 0 on a clean tree" || bad "check exit=$rc on clean tree"
|
||||
python3 - "$WORK/report.json" <<'PY' && ok "check report is valid JSON with ok=true and empty diffs" || bad "report shape wrong"
|
||||
import json, sys
|
||||
r = json.load(open(sys.argv[1]))
|
||||
assert r["ok"] is True, r
|
||||
assert r["missing"] == [] and r["changed"] == [] and r["added"] == [] and r["stale"] == [], r
|
||||
assert r["locked_tree_digest"] == r["observed_tree_digest"], r
|
||||
assert r["file_count"] > 0, r
|
||||
PY
|
||||
|
||||
printf '\n%d failure(s)\n' "$fails"
|
||||
[ "$fails" -eq 0 ]
|
||||
Reference in new issue
Block a user