diff --git a/src/hbkit/archive.py b/src/hbkit/archive.py index 64d8a71..29ace13 100644 --- a/src/hbkit/archive.py +++ b/src/hbkit/archive.py @@ -136,12 +136,13 @@ def __init__(self, d: str): found[n] = (gen, f) if not found: raise FileNotFoundError(f"no .idx shards in {d}") - self.files = [found[k][1] for k in sorted(found)] - n = len(self.files) + shard_ids = sorted(found) + self.files = [found[k][1] for k in shard_ids] last = os.path.getsize(os.path.join(d, self.files[-1])) - self.sizes = [SHARD_SIZE] * (n - 1) + [last] - self.starts = [i * SHARD_SIZE for i in range(n)] - self.total = (n - 1) * SHARD_SIZE + last + self.sizes = [SHARD_SIZE] * (len(self.files) - 1) + [last] + self.starts = [k * SHARD_SIZE for k in shard_ids] + self.total = self.starts[-1] + last + self._positions = dict(zip(shard_ids, range(len(shard_ids)))) self._fh: dict[int, object] = {} # Shard 0 carries a 64-byte header that declares the record size, so readers do # not have to hardcode per-DSM-version layouts. @@ -160,15 +161,18 @@ def read(self, off: int, n: int) -> bytes: """Read n bytes at logical offset off, spanning shards as needed.""" out = b"" i = off // SHARD_SIZE - while n > 0 and i < len(self.files): - st = self.starts[i] - take = min(n, st + self.sizes[i] - off) + while n > 0: + pos = self._positions.get(i) + if pos is None: + break + st = self.starts[pos] + take = min(n, st + self.sizes[pos] - off) if take <= 0: break - fh = self._fh.get(i) + fh = self._fh.get(pos) if fh is None: - fh = self._fh[i] = open(os.path.join(self.d, self.files[i]), "rb", - buffering=INDEX_BUF) + fh = self._fh[pos] = open(os.path.join(self.d, self.files[pos]), "rb", + buffering=INDEX_BUF) fh.seek(off - st) got = fh.read(take) out += got diff --git a/src/hbkit/cli.py b/src/hbkit/cli.py index 8b5c10f..14bfab9 100644 --- a/src/hbkit/cli.py +++ b/src/hbkit/cli.py @@ -231,7 +231,23 @@ def main() -> int: if cmd == "doctor": from . import doctor - r = doctor.diagnose(archive, password=resolve_password(password)) + sample = None + for flag in ("-s", "-n", "--sample"): + if flag in a[2:]: + i = a.index(flag, 2) + try: + sample = int(a[i + 1]) + except (IndexError, ValueError): + print(f"{flag} requires a positive integer", file=sys.stderr) + return 2 + if sample < 1: + print(f"{flag} requires a positive integer", file=sys.stderr) + return 2 + break + kw = {"password": resolve_password(password)} + if sample is not None: + kw["sample"] = sample + r = doctor.diagnose(archive, **kw) print(doctor.render(r)) return 0 if (r.ok and not r.blockers) else 1 diff --git a/src/hbkit/doctor.py b/src/hbkit/doctor.py index 4f4d5cd..161e66a 100644 --- a/src/hbkit/doctor.py +++ b/src/hbkit/doctor.py @@ -202,7 +202,7 @@ def main() -> int: ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("archive") ap.add_argument("-p", "--password", help="for encrypted archives") - ap.add_argument("-n", "--sample", type=int, default=10, + ap.add_argument("-n", "-s", "--sample", type=int, default=10, help="files to rebuild per share as proof (default 10)") a = ap.parse_args() r = diagnose(a.archive, a.sample, password=a.password) diff --git a/tests/test_hbkit.py b/tests/test_hbkit.py index a3003d8..d05ca12 100644 --- a/tests/test_hbkit.py +++ b/tests/test_hbkit.py @@ -287,6 +287,18 @@ def test_cat_single_shard_still_works(tmp_path, monkeypatch): c.close() +def test_cat_preserves_sparse_shard_offsets(tmp_path, monkeypatch): + from hbkit import archive as A + blob = _make_shards(tmp_path, [128, 10], 128, monkeypatch) + (tmp_path / "1.idx.2").rename(tmp_path / "2.idx.2") + c = A.Cat(str(tmp_path)) + assert c.starts == [0, 256] + assert c.total == 266 + assert c.read(128, 1) == b"" + assert c.read(256, 10) == blob[128:] + c.close() + + # --------------------------------------------------- rclone mount helper (no network) def test_mount_profiles_pick_the_right_cache_mode(monkeypatch): @@ -480,6 +492,24 @@ class Result: assert seen["pw"] == want +@pytest.mark.parametrize("flag", ["-s", "-n", "--sample"]) +def test_doctor_forwards_sample_count(monkeypatch, flag): + from hbkit import cli, doctor + + seen = {} + + class Result: + ok, blockers = True, [] + + monkeypatch.setattr(doctor, "diagnose", + lambda a, sample=10, password=None: + (seen.update(sample=sample), Result())[1]) + monkeypatch.setattr(doctor, "render", lambda r: "") + monkeypatch.setattr(sys, "argv", ["hbk", "/some/archive", "doctor", flag, "2"]) + assert cli.main() == 0 + assert seen["sample"] == 2 + + def test_lazycats_lists_shards_before_any_are_opened(): """`doctor` describes an archive without extracting from it, so nothing has opened a file_chunk shard by the time it prints them. Iterating the dict reported an empty list