diff --git a/tests/test_bill_index.py b/tests/test_bill_index.py index 4a715e76..f7f73f02 100644 --- a/tests/test_bill_index.py +++ b/tests/test_bill_index.py @@ -28,7 +28,7 @@ import pytest -from bill_index.bill_index import ( +from tools.shared.bill_index import ( BillIdentifier, BillIndex, _decode_value, diff --git a/tests/test_fetch_bill_archives.py b/tests/test_fetch_bill_archives.py index 39f5f677..bcefc85f 100644 --- a/tests/test_fetch_bill_archives.py +++ b/tests/test_fetch_bill_archives.py @@ -15,7 +15,8 @@ import pytest import respx -from fetch_bill_archives import archive_temp_path, download_archive_zip +from tests.utils import assert_files +from tools.shared.http import download_zip as download_archive_zip ARCHIVE_URL = "https://www.govinfo.gov/bulkdata/BILLSTATUS/999/hr/BILLSTATUS-999-hr.zip" @@ -53,8 +54,7 @@ def test_truncated_body_without_content_length_is_not_committed(self, tmp_path): with pytest.raises(httpx.HTTPError): download_archive_zip(client, ARCHIVE_URL, dest) - assert not dest.exists() - assert not archive_temp_path(dest).exists() + assert_files(tmp_path, ["999-hr.zip.error"]) @respx.mock def test_healthy_body_without_content_length_is_committed(self, tmp_path): @@ -67,7 +67,7 @@ def test_healthy_body_without_content_length_is_committed(self, tmp_path): download_archive_zip(client, ARCHIVE_URL, dest) assert dest.read_bytes() == full - assert not archive_temp_path(dest).exists() + assert_files(tmp_path, ["999-hr.zip"]) with zipfile.ZipFile(dest) as zf: assert zf.namelist() == ["BILLSTATUS-999hr1.xml"] diff --git a/tests/test_fetch_bill_archives_extract.py b/tests/test_fetch_bill_archives_extract.py index c77a840c..e88ce486 100644 --- a/tests/test_fetch_bill_archives_extract.py +++ b/tests/test_fetch_bill_archives_extract.py @@ -16,182 +16,38 @@ from __future__ import annotations -import io -import struct -import time -import zipfile +import shlex from pathlib import Path +from unittest.mock import patch -import httpx import pytest import respx -import fetch_bill_archives -from fetch_bill_archives import ( - MAX_MEMBER_COUNT, - MAX_UNCOMPRESSED_BYTES, - archive_destination, - archive_error_path, - archive_extract_dir, - archive_url, - download_archives, - extract_archive, - extract_archives, +import tools.shared.http as http +from tests.utils import archive_bytes, assert_files, mock_http_requests +from tools.fetch_bill_archives import ( + billstatus_filename, + billstatus_zip_filename, + main, +) +from tools.fetch_bill_archives import ( + billstatus_zip_url as archive_url, ) -def write_archive(source: Path, name: str, members: dict[str, bytes] | None = None) -> Path: - """Write a well-formed ZIP named ``{name}.zip`` into source.""" - # `is None`, not `or`: an explicitly empty dict means a zero-member ZIP, and - # falling back on falsiness would silently write a one-member archive instead. - if members is None: - members = {f"{name}-1.xml": b""} - path = source / f"{name}.zip" - buf = io.BytesIO() - with zipfile.ZipFile(buf, "w") as zf: - for member, body in members.items(): - zf.writestr(member, body) - path.write_bytes(buf.getvalue()) - return path - - -def write_oversized_archive(source: Path, name: str, declared_sizes: list[int]) -> Path: - """A ZIP whose central directory DECLARES ``declared_sizes`` uncompressed bytes - per member while actually holding a few. - - A decompression bomb has exactly this shape at scale: a small archive whose - members expand enormously on disk. Materializing a real multi-gigabyte expansion - in a test would spend the disk the ceiling exists to protect, so the declaration - is patched instead -- and the declaration is what the ceiling reads, because - ``zipfile`` takes member sizes from the central directory. - """ - path = source / f"{name}.zip" - buf = io.BytesIO() - with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf: - for i in range(len(declared_sizes)): - zf.writestr(f"{name}-{i}.xml", b"") - raw = bytearray(buf.getvalue()) - # Central directory file header: uncompressed size is a 4-byte LE field at +24. - offset = 0 - for declared in declared_sizes: - offset = raw.find(b"PK\x01\x02", offset) - assert offset != -1, "central directory record not found -- the craft is broken" - struct.pack_into(" bytes: - """One well-formed BILLSTATUS archive ZIP, as bytes. - - A pure function of ``name``: same argument, same bytes, whenever it is called. - That is what makes it usable on both sides of a byte-equality assertion, which - ``TestArchiveBytesIsDeterministic`` pins and the cached-archive test relies on. - Passing an explicit ``ZipInfo`` is what buys it -- ``writestr`` given a bare - arcname stamps the member with the wall clock instead (#480). - """ - info = zipfile.ZipInfo(f"{name}-1.xml", FIXED_ZIP_DATE_TIME) - info.external_attr = NORMAL_FILE_EXTERNAL_ATTR - buf = io.BytesIO() - with zipfile.ZipFile(buf, "w") as zf: - zf.writestr(info, b"") - return buf.getvalue() - - -class TestArchiveBytesIsDeterministic: - """The fixture's own contract, because a test below leans on it (#480). - - ``test_a_cached_archive_clears_a_stale_marker`` proves the cached archive was not - rewritten, by comparing the file on disk against a freshly generated one. That - comparison only means "unchanged" if identical content implies identical bytes. - It did not: ``writestr`` given a bare arcname stamps the member with - ``time.localtime()``, which the DOS timestamp field carries at two-second - granularity -- so the assertion held only while both generations landed in the - same tick, and failed on a loaded run. The failure surfaced as a cache-coherence - bug in ``download_archives``, which is the expensive part: the two byte strings - print identically in the truncated repr, so nothing about the output points at - the fixture. - """ +def fetch_bill_archives(args: str) -> None: + argv = ["fetch_bill_archives", *shlex.split(args)] + with patch("sys.argv", argv): + return main() - def test_output_does_not_depend_on_the_wall_clock(self, monkeypatch): - def at(clock, build): - with monkeypatch.context() as m: - # A no-op through 3.13, where `writestr` reads `time.localtime` - # unconditionally. From 3.14 it resolves the stamp through - # `ZipInfo._for_archive`, which prefers SOURCE_DATE_EPOCH via - # `time.gmtime` -- so with that variable set in the environment the - # control below would see one timestamp under both clocks and fail, - # reporting a broken gate rather than a set env var. Version boundary - # checked against real 3.13.15 and 3.14.6 interpreters, not inferred. - m.delenv("SOURCE_DATE_EPOCH", raising=False) - m.setattr(time, "localtime", lambda *a, **k: time.struct_time(clock)) - return build() - - early = (1999, 12, 31, 23, 59, 58, 4, 365, 0) - late = (2044, 7, 4, 12, 30, 20, 0, 186, 0) - - # The clock is moved rather than waited on, because generating twice in quick - # succession passes against the broken helper too -- both calls share a tick, - # so that check could never fail and would certify nothing. - # - # This control does double duty. It is the defect held next to the fix, and it - # is the only thing proving the patch above actually reaches zipfile: if the - # clock were not really moving, the assertion below would pass however the - # helper were written. - def stamped_by_the_wall_clock() -> bytes: - buf = io.BytesIO() - with zipfile.ZipFile(buf, "w") as zf: - zf.writestr("119-hr-1.xml", b"") - return buf.getvalue() - - assert at(early, stamped_by_the_wall_clock) != at(late, stamped_by_the_wall_clock) - - assert at(early, archive_bytes) == at(late, archive_bytes) - - def test_only_the_timestamp_differs_from_the_old_fixture(self, monkeypatch): - # Supplying an explicit ``ZipInfo`` opts out of every default ``writestr`` would - # have filled in, not just the timestamp -- compression, flags, and the member's - # mode among them. Pinning the OLD construction to the same timestamp isolates - # that: with the clock held equal, any surviving difference is metadata the - # repair moved by accident rather than the clock it moved on purpose. - # - # Compared over raw bytes rather than an enumerated field list, because the - # field nobody thought to enumerate is exactly the one that would slip through. - def old_construction() -> bytes: - buf = io.BytesIO() - with zipfile.ZipFile(buf, "w") as zf: - zf.writestr("119-hr-1.xml", b"") - return buf.getvalue() - monkeypatch.delenv("SOURCE_DATE_EPOCH", raising=False) - monkeypatch.setattr(time, "localtime", lambda *a, **k: time.struct_time(FIXED_ZIP_DATE_TIME + (0, 1, 0))) +def single_member_archive_bytes(congress: int, bill_type: str) -> bytes: + return archive_bytes({billstatus_filename(congress, bill_type, 1): b""}) - assert old_construction() == archive_bytes() - def test_the_fixed_timestamp_did_not_cost_the_archive(self): - # A helper that returned a constant would satisfy the test above perfectly, so - # the output still has to be a real archive -- otherwise the determinism gate - # could be met by breaking every test that consumes the fixture. - with zipfile.ZipFile(io.BytesIO(archive_bytes())) as zf: - assert zf.namelist() == ["119-hr-1.xml"] - assert zf.read("119-hr-1.xml") == b"" - assert zf.infolist()[0].date_time == FIXED_ZIP_DATE_TIME +def write_single_member_archive(source: Path, congress: int, bill_type: str) -> Path: + path = source / billstatus_zip_filename(congress, bill_type) + path.write_bytes(single_member_archive_bytes(congress, bill_type)) class TestDownloadArchivesCacheCoherence: @@ -207,68 +63,74 @@ def test_existing_archive_is_skipped_without_a_request(self, tmp_path): # re-downloads hundreds of MB. Asserting the route was never called is the # point -- an implementation that fetched and then discarded the body would # leave the same files on disk. - dest = archive_destination(tmp_path, 119, "hr") - dest.write_bytes(b"cached, not a real zip") - route = respx.get(archive_url(119, "hr")).mock(return_value=httpx.Response(200)) + write_single_member_archive(tmp_path, 119, "hr") + route = mock_http_requests(archive_url(119, "hr"), content=b"new content") - downloaded = download_archives(119, 119, bill_types=["hr"], destination=tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - assert downloaded == [] # skipped, so not reported as newly downloaded assert not route.called - assert dest.read_bytes() == b"cached, not a real zip" @respx.mock def test_successful_download_is_committed_and_reported(self, tmp_path): - body = archive_bytes() - respx.get(archive_url(119, "hr")).mock(return_value=httpx.Response(200, content=body)) + body = single_member_archive_bytes(119, "hr") + mock_http_requests(archive_url(119, "hr"), content=body) - downloaded = download_archives(119, 119, bill_types=["hr"], destination=tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - dest = archive_destination(tmp_path, 119, "hr") - assert downloaded == [dest] - assert dest.read_bytes() == body - assert not archive_error_path(tmp_path, 119, "hr").exists() + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip", "119-hr-1"]) + assert_files(tmp_path / "119-hr-1", ["119-hr-1_status.xml"]) @respx.mock def test_failed_download_writes_an_error_marker_and_commits_no_archive(self, tmp_path): - respx.get(archive_url(119, "hr")).mock(return_value=httpx.Response(404)) + mock_http_requests(archive_url(119, "hr"), status_code=404) - downloaded = download_archives(119, 119, bill_types=["hr"], destination=tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - assert downloaded == [] - # No .zip: a committed-but-failed archive would be skipped forever by the - # test above, permanently caching the failure as if it were data. - assert not archive_destination(tmp_path, 119, "hr").exists() - error_path = archive_error_path(tmp_path, 119, "hr") - assert error_path.exists() - assert error_path.read_text(encoding="utf-8") # carries the reason, not empty + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip.error"]) @respx.mock def test_a_later_success_clears_a_stale_error_marker(self, tmp_path): # The coherence half: without the clear, a marker from a transient outage # would keep describing a failure for an archive that is now present, and # anything reading markers to decide what is missing would be wrong forever. - error_path = archive_error_path(tmp_path, 119, "hr") - error_path.parent.mkdir(parents=True, exist_ok=True) + error_path = tmp_path / "BILLSTATUS-119-hr.zip.error" error_path.write_text("earlier failure", encoding="utf-8") - respx.get(archive_url(119, "hr")).mock(return_value=httpx.Response(200, content=archive_bytes())) + mock_http_requests(archive_url(119, "hr"), content=single_member_archive_bytes(119, "hr")) + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip.error"]) - download_archives(119, 119, bill_types=["hr"], destination=tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - assert not error_path.exists() - assert archive_destination(tmp_path, 119, "hr").exists() + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip", "119-hr-1"]) + assert_files(tmp_path / "119-hr-1", ["119-hr-1_status.xml"]) @respx.mock def test_a_failing_archive_does_not_abort_the_batch(self, tmp_path): - # Tasks are newest-first, so 119 is attempted before 118. - respx.get(archive_url(119, "hr")).mock(return_value=httpx.Response(500)) - respx.get(archive_url(118, "hr")).mock(return_value=httpx.Response(200, content=archive_bytes("118-hr"))) - - downloaded = download_archives(118, 119, bill_types=["hr"], destination=tmp_path) - - assert downloaded == [archive_destination(tmp_path, 118, "hr")] - assert archive_error_path(tmp_path, 119, "hr").exists() - assert not archive_error_path(tmp_path, 118, "hr").exists() + mock_http_requests(archive_url(117, "hr"), status_code=500) + mock_http_requests(archive_url(118, "hr"), content=single_member_archive_bytes(118, "hr")) + mock_http_requests(archive_url(119, "hr"), status_code=500) + + fetch_bill_archives( + f"--from-congress 117 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) + + assert_files( + tmp_path, + { + "BILLSTATUS-117-hr.zip.error", + "BILLSTATUS-118-hr.zip", + "BILLSTATUS-119-hr.zip.error", + "118-hr-1", + }, + ) + assert_files(tmp_path / "118-hr-1", ["118-hr-1_status.xml"]) @respx.mock def test_a_stale_error_marker_alone_does_not_prevent_a_retry(self, tmp_path): @@ -276,14 +138,19 @@ def test_a_stale_error_marker_alone_does_not_prevent_a_retry(self, tmp_path): # attempted; only a present .zip suppresses the request. Pinning this keeps a # future "skip anything with an .error marker" optimization from silently # making transient failures permanent. - error_path = archive_error_path(tmp_path, 119, "hr") - error_path.parent.mkdir(parents=True, exist_ok=True) + target_archive: Path = tmp_path / "BILLSTATUS-119-hr.zip" + error_path = http.path_for_error(target_archive) error_path.write_text("earlier failure", encoding="utf-8") - route = respx.get(archive_url(119, "hr")).mock(return_value=httpx.Response(200, content=archive_bytes())) + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip.error"]) - download_archives(119, 119, bill_types=["hr"], destination=tmp_path) + route = mock_http_requests(archive_url(119, "hr"), content=single_member_archive_bytes(119, "hr")) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) assert route.called + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip", "119-hr-1"]) + assert_files(tmp_path / "119-hr-1", ["119-hr-1_status.xml"]) @respx.mock def test_a_cached_archive_clears_a_stale_marker(self, tmp_path): @@ -292,37 +159,43 @@ def test_a_cached_archive_clears_a_stale_marker(self, tmp_path): # describing a failure that the archive itself disproves. The two states # together are contradictory, so whichever a future consumer reads, it reads a # wrong answer for one of them. - dest = archive_destination(tmp_path, 119, "hr") - dest.write_bytes(archive_bytes()) - error_path = archive_error_path(tmp_path, 119, "hr") + dest = tmp_path / "BILLSTATUS-119-hr.zip" + dest.write_bytes(single_member_archive_bytes(119, "hr")) + error_path = tmp_path / "BILLSTATUS-119-hr.zip.error" error_path.write_text("earlier failure", encoding="utf-8") - downloaded = download_archives(119, 119, bill_types=["hr"], destination=tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - assert not error_path.exists() - # Still skipped, not re-fetched: clearing the marker must not cost the cache. - # respx is mocking with no route registered, so any request would raise. - assert downloaded == [] - assert dest.read_bytes() == archive_bytes() + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip", "119-hr-1"]) + assert_files(tmp_path / "119-hr-1", ["119-hr-1_status.xml"]) class TestExtractArchive: + @respx.mock def test_creates_the_destination_and_writes_members(self, tmp_path): - archive = write_archive(tmp_path, "119-hr", {"a.xml": b"", "sub/b.xml": b""}) + archive = tmp_path / "BILLSTATUS-119-hr.zip" + archive.write_bytes(single_member_archive_bytes(119, "hr")) dest = tmp_path / "out" - extract_archive(archive, dest) + fetch_bill_archives(f"--from-congress 119 --to-congress 119 --types hr --out-dir {dest} --zip-dir {tmp_path}") - assert (dest / "a.xml").read_bytes() == b"" - assert (dest / "sub" / "b.xml").read_bytes() == b"" + assert_files(dest, ["119-hr-1"]) + assert_files(dest / "119-hr-1", ["119-hr-1_status.xml"]) @pytest.mark.parametrize( "member", [ - pytest.param("../escaped.xml", id="parent-traversal"), - pytest.param("../../escaped.xml", id="double-parent-traversal"), - pytest.param("/etc/escaped.xml", id="absolute-path"), - pytest.param("sub/../../escaped.xml", id="traversal-after-descent"), + # Relative paths + pytest.param("../escaped.xml", id="relative-parent-traversal"), + pytest.param("../../escaped.xml", id="relative-double-parent-traversal"), + pytest.param("sub/../../escaped.xml", id="relative-traversal-after-descent"), + pytest.param("../sub/escaped.xml", id="relative-parent-then-descent"), + # Absolute paths + pytest.param("/etc/escaped.xml", id="absolute-unix"), + pytest.param("/tmp/escaped.xml", id="absolute-tmp"), + pytest.param("//escaped.xml", id="absolute-double-slash"), ], ) def test_members_cannot_escape_the_destination_directory(self, tmp_path, member): @@ -339,250 +212,136 @@ def test_members_cannot_escape_the_destination_directory(self, tmp_path, member) # on global filesystem state -- shared with every other process on the # machine, so a stale file from an unrelated run fails it and a concurrent # run makes it flaky. - archive = write_archive(tmp_path, "119-hr", {member: b""}) - dest = tmp_path / "out" - - extract_archive(archive, dest) - - written = [p for p in dest.rglob("*") if p.is_file()] - assert written, "nothing extracted, so the containment check proved nothing" - for path in written: - assert dest.resolve() in path.resolve().parents - # Nothing escaped one level up into the directory holding the archive. - assert not (tmp_path / "escaped.xml").exists() - - def test_archive_declaring_more_than_the_ceiling_is_refused(self, tmp_path): - # #279: zipfile.extractall applies no bound, so a crafted archive could fill - # the disk. The refusal must land BEFORE anything is written -- a ceiling - # checked while extracting would already have spent the disk it protects. - archive = write_oversized_archive(tmp_path, "119-hr", [MAX_UNCOMPRESSED_BYTES + 1]) - dest = tmp_path / "out" - - with pytest.raises(ValueError, match="uncompressed"): - extract_archive(archive, dest) - - assert not dest.exists() - - def test_the_ceiling_counts_every_member_not_just_the_largest(self, tmp_path): - # A bomb split across members would slip a per-member check: each part sits - # comfortably under the ceiling and only the total is ruinous. Four members at - # 40% each are individually fine and collectively 1.6x over. - part = int(MAX_UNCOMPRESSED_BYTES * 0.4) - archive = write_oversized_archive(tmp_path, "119-hr", [part] * 4) - dest = tmp_path / "out" - - with pytest.raises(ValueError, match="uncompressed"): - extract_archive(archive, dest) - - assert not dest.exists() - - def test_an_archive_at_the_ceiling_is_still_extracted(self, tmp_path): - # The boundary is inclusive, and more to the point the ceiling must not be so - # eager that it refuses work: a real archive is ~162 MiB expanded, three - # orders of magnitude under this, and must extract untouched. A test that only - # ever saw the refusal could not tell a working ceiling from one wired to - # reject everything. - archive = write_oversized_archive(tmp_path, "119-hr", [MAX_UNCOMPRESSED_BYTES]) - dest = tmp_path / "out" - - extract_archive(archive, dest) - - assert (dest / "119-hr-0.xml").read_bytes() == b"" - - def test_archive_with_more_members_than_the_ceiling_is_refused(self, tmp_path, monkeypatch): - # #306: the byte ceiling sums declared uncompressed sizes, so an archive of - # empty members declares zero bytes and sails past it no matter how many - # members it holds -- 200,000 empty files weigh nothing yet exhaust inodes. - # The member ceiling is what catches it. Members are empty on purpose (declared - # size 0), so this archive passes the byte ceiling and only the count refuses - # it, proving the gap the issue describes is actually closed. The ceiling is - # lowered rather than materializing 100,001 members, whose sole cost is build - # time -- the check reads len(infolist()), which needs that many real central - # directory records, so the real bound cannot be faked cheaply the way the byte - # bomb's declared sizes can. - monkeypatch.setattr(fetch_bill_archives, "MAX_MEMBER_COUNT", 3) - archive = write_archive(tmp_path, "119-hr", {f"119-hr-{i}.xml": b"" for i in range(4)}) - dest = tmp_path / "out" - - with pytest.raises(ValueError, match="4 members, over the 3 ceiling"): - extract_archive(archive, dest) - - assert not dest.exists() - - def test_an_archive_at_the_member_ceiling_is_still_extracted(self, tmp_path, monkeypatch): - # The boundary is inclusive, and the point mirrors the byte ceiling's: a bound - # wired to reject everything would pass a refusal test while breaking real - # extraction. A real archive is ~10,564 members, an order of magnitude under - # the true ceiling, and must extract untouched. - monkeypatch.setattr(fetch_bill_archives, "MAX_MEMBER_COUNT", 3) - archive = write_archive(tmp_path, "119-hr", {f"119-hr-{i}.xml": b"" for i in range(3)}) dest = tmp_path / "out" + dest.mkdir() + archive_dest = dest / "BILLSTATUS-119-hr.zip" + bytes = archive_bytes({member: b""}) + archive_dest.write_bytes(bytes) + route = mock_http_requests(archive_url(119, "hr"), content=bytes) - extract_archive(archive, dest) + fetch_bill_archives(f"--from-congress 119 --to-congress 119 --types hr --out-dir {dest} --zip-dir {dest}") + assert not route.called - assert (dest / "119-hr-0.xml").read_bytes() == b"" - assert len([p for p in dest.rglob("*") if p.is_file()]) == 3 + assert_files(tmp_path, ["out"]) + assert_files(dest, {"BILLSTATUS-119-hr.zip"}) - def test_member_ceiling_calibration_excludes_the_306_exhaustion_case(self): - # Two-sided calibration guard (#447): a floor alone only stops the ceiling - # being lowered too far -- raising it to 50,000,000, or effectively off at - # 10**12, passed the full suite before this upper bound existed. - # - # The floor is the largest known real BILLSTATUS archive, 10,564 members - # (#300). The ceiling is pinned to the CURRENT production value, 100,000 -- - # NOT merely kept under 200,000, the member count #306 used to demonstrate - # inode exhaustion. `extract_archive` refuses only on - # `member_count > MAX_MEMBER_COUNT`, so any ceiling below 200,000, even - # 199,999, technically still refuses that exact archive, but only after - # nearly as many inodes are counted as the demonstrated-bad case consumed -- - # eroding the safety margin to nothing while this assertion stayed green. - # The production comment's own stated rationale for 100,000 is a 2x margin - # under that threshold ("refusing the 200,000-member archive... by a factor - # of two"), so the calibration pins that margin explicitly rather than - # tolerating any widening up to one member short of #306's demonstration. - # - # No companion "still extracts a normal archive" test here: that positive - # control already exists and doesn't need duplicating -- - # test_creates_the_destination_and_writes_members above extracts a small - # archive under both real, unpatched ceilings. - assert 10_564 < MAX_MEMBER_COUNT <= 100_000 - - def test_byte_ceiling_calibration_has_real_corpus_headroom(self): - # Companion to the member-ceiling calibration above, for MAX_UNCOMPRESSED_BYTES - # (#447): before this test, the byte ceiling had no calibration assertion at - # all, one-sided or otherwise -- every byte-ceiling test builds its fixture - # FROM the constant (write_oversized_archive(..., [MAX_UNCOMPRESSED_BYTES + 1])), - # so the fixture rescales with whatever the constant is and can never disagree - # with it. Widening the constant to 3,900,000,000 -- nearly double its real - # value -- passed the full suite. - # - # The floor is the largest known real archive's expanded size, ~162 MiB - # (#300). Unlike the member ceiling, there is no single demonstrated-bad byte - # figure in the tracker to exclude, so the ceiling is set from the constant's - # own documented margin claim ("2 GiB leaves better than an order of - # magnitude of headroom" over 162 MiB): 3 GiB is roughly 19x the floor, - # comfortably inside that claimed order of magnitude while still catching a - # ceiling widened toward "off". - # - # No companion extraction test here either, for the same reason as the member - # ceiling above: test_creates_the_destination_and_writes_members and - # test_an_archive_at_the_ceiling_is_still_extracted already cover normal - # extraction under the real, unpatched byte ceiling. - assert 162 * 1024**2 < MAX_UNCOMPRESSED_BYTES <= 3 * 1024**3 - - def test_raises_on_a_corrupt_archive(self, tmp_path): - archive = tmp_path / "119-hr.zip" + @respx.mock + def test_does_not_extract_on_a_corrupt_archive(self, tmp_path): + archive = tmp_path / "BILLSTATUS-119-hr.zip" archive.write_bytes(b"not a zip") - with pytest.raises(zipfile.BadZipFile): - extract_archive(archive, tmp_path / "out") + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) + + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip"]) class TestExtractArchivesCacheCoherence: + @respx.mock def test_extracts_each_archive_into_a_folder_named_for_its_stem(self, tmp_path): # Written out of alphabetical order so the assertion below tests the sort # rather than the order the files happened to be created in. - write_archive(tmp_path, "119-s") - write_archive(tmp_path, "119-hr") + write_single_member_archive(tmp_path, 119, "s") + write_single_member_archive(tmp_path, 119, "hr") + out_dir = tmp_path / "out" - extracted = extract_archives(tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr s --out-dir {out_dir} --zip-dir {tmp_path}" + ) - # Compared in order, not sorted: extract_archives sorts its glob, and other - # tests here rely on that ordering being deterministic (the batch-resilience - # test needs the corrupt archive to come first). Sorting the actual value - # would discard the very property those tests lean on, leaving the return - # order free to follow filesystem order unnoticed. - assert [p.name for p in extracted] == ["119-hr", "119-s"] - assert (tmp_path / "119-hr" / "119-hr-1.xml").exists() + assert_files(out_dir, ["119-hr-1", "119-s-1"]) + @respx.mock def test_existing_folder_is_skipped_and_left_untouched(self, tmp_path): # The skip is keyed on folder existence alone, so a pre-existing folder wins # over the archive's actual contents. Pinning it keeps the re-run cheap and # documents that the folder, not the ZIP, is the cache. - write_archive(tmp_path, "119-hr", {"fresh.xml": b""}) - stale_dir = tmp_path / "119-hr" + write_single_member_archive(tmp_path, 119, "hr") + stale_dir = tmp_path / "119-hr-1" stale_dir.mkdir() - (stale_dir / "stale.xml").write_bytes(b"") + (stale_dir / "119-hr-1_status.xml").write_bytes(b"") - extracted = extract_archives(tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - assert extracted == [] # skipped, so not reported as newly extracted - assert (stale_dir / "stale.xml").exists() - assert not (stale_dir / "fresh.xml").exists() + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip", "119-hr-1"]) + assert_files(stale_dir, ["119-hr-1_status.xml"]) + assert (stale_dir / "119-hr-1_status.xml").read_bytes() == b"" + @respx.mock def test_partial_folder_from_a_failed_extract_is_removed(self, tmp_path): # The cleanup is what makes the existence-based skip safe: a folder left # behind here would be treated as a complete extraction by every later run, # silently serving a truncated corpus. - corrupt = tmp_path / "119-hr.zip" + corrupt = tmp_path / "BILLSTATUS-119-hr.zip" corrupt.write_bytes(b"not a zip") - extracted = extract_archives(tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - assert extracted == [] - assert not archive_extract_dir(tmp_path, corrupt).exists() + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip"]) + @respx.mock def test_a_failed_archive_does_not_abort_the_batch(self, tmp_path): # Sorted order puts the corrupt archive first, so a bare raise would cost the # healthy ones too. - (tmp_path / "119-aaa.zip").write_bytes(b"not a zip") - write_archive(tmp_path, "119-zzz") + (tmp_path / "BILLSTATUS-119-aaa.zip").write_bytes(b"not a zip") + write_single_member_archive(tmp_path, 119, "zzz") - extracted = extract_archives(tmp_path) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types aaa zzz --out-dir {tmp_path} --zip-dir {tmp_path}" + ) - assert [p.name for p in extracted] == ["119-zzz"] - assert (tmp_path / "119-zzz" / "119-zzz-1.xml").exists() - assert not (tmp_path / "119-aaa").exists() + assert_files(tmp_path, ["BILLSTATUS-119-aaa.zip", "BILLSTATUS-119-zzz.zip", "119-zzz-1"]) + assert_files(tmp_path / "119-zzz-1", ["119-zzz-1_status.xml"]) + @respx.mock def test_only_zip_files_are_considered(self, tmp_path, capsys): # The bills directory holds bills.csv and extracted folders alongside the # archives, so the glob is what keeps them out. Asserting the extracted list # alone would not catch a widened glob: a non-ZIP that gets attempted fails to # open and is swallowed by the same except that handles a corrupt archive, so # the list comes out identical either way and only the log betrays it. - write_archive(tmp_path, "119-hr") + write_single_member_archive(tmp_path, 119, "hr") (tmp_path / "notes.txt").write_text("ignore me") (tmp_path / "bills.csv").write_text("id\n") - extracted = extract_archives(tmp_path) - - assert [p.name for p in extracted] == ["119-hr"] - assert not (tmp_path / "notes").exists() + mock_http_requests(archive_url(119, "hr"), content=single_member_archive_bytes(119, "hr")) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) err = capsys.readouterr().err assert "notes.txt" not in err assert "bills.csv" not in err - def test_zero_member_archive_still_creates_its_folder(self, tmp_path): - # A zero-member ZIP is structurally valid and is deliberately not treated as a - # failed download (see _verify_archive_complete). zipfile.extractall does not - # create the destination when there is nothing to write, so the explicit - # mkdir in extract_archive is the only thing that does -- and without the - # folder the run would report an extraction that left no cache entry, so the - # next run would extract it again instead of skipping. - write_archive(tmp_path, "119-hr", members={}) - - extracted = extract_archives(tmp_path) + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip", "119-hr-1", "notes.txt", "bills.csv"]) - assert [p.name for p in extracted] == ["119-hr"] - assert (tmp_path / "119-hr").is_dir() - assert extract_archives(tmp_path) == [] # now a coherent cache entry - - def test_empty_source_directory_yields_nothing(self, tmp_path): - assert extract_archives(tmp_path) == [] + @respx.mock + def test_zero_member_archive_writes_no_bill_folder(self, tmp_path): + # A zero-member ZIP is structurally valid and is deliberately not treated as a + # failed download (see _verify_archive_complete). Archives hold many bills, so + # an empty zip has no bill id to name a folder after -- extraction is a no-op. + (tmp_path / billstatus_zip_filename(119, "hr")).write_bytes(archive_bytes({})) - def test_missing_source_directory_raises(self, tmp_path): - missing = tmp_path / "nope" - with pytest.raises(ValueError, match="Source folder does not exist"): - extract_archives(missing) + fetch_bill_archives( + f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + ) + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip"]) + @respx.mock def test_rerun_after_a_successful_extract_is_a_no_op(self, tmp_path): - # The cache-coherence property stated end to end: extract, then extract again - # and get nothing new, with the first run's output intact. - write_archive(tmp_path, "119-hr") - - first = extract_archives(tmp_path) - second = extract_archives(tmp_path) - - assert [p.name for p in first] == ["119-hr"] - assert second == [] - assert (tmp_path / "119-hr" / "119-hr-1.xml").exists() + zip_name = "BILLSTATUS-119-hr.zip" + member = "BILLSTATUS-119hr1.xml" + (tmp_path / zip_name).write_bytes(archive_bytes({member: b"first contents"})) + args = f"--from-congress 119 --to-congress 119 --types hr --out-dir {tmp_path} --zip-dir {tmp_path}" + fetch_bill_archives(args) + + (tmp_path / zip_name).write_bytes(archive_bytes({member: b"replacement contents"})) + fetch_bill_archives(args) + + assert_files(tmp_path, ["BILLSTATUS-119-hr.zip", "119-hr-1"]) + assert_files(tmp_path / "119-hr-1", ["119-hr-1_status.xml"]) + assert (tmp_path / "119-hr-1" / "119-hr-1_status.xml").read_bytes() == b"first contents" diff --git a/tests/test_fetch_bills.py b/tests/test_fetch_bills.py index c4fc168d..cbedabd9 100644 --- a/tests/test_fetch_bills.py +++ b/tests/test_fetch_bills.py @@ -3,17 +3,20 @@ import argparse import json import time +import zipfile import httpx import pytest import respx +import fetch_bills as fb import fetch_govinfo as gi from fetch_bills import ( api_get, build_parser, cmd_download, cmd_download_all, + cmd_fetch_index, congress_for_year, download_all_versions, download_version_xml, @@ -27,10 +30,15 @@ save_version, version_path, ) +from tests.utils import EMPTY_ZIP_BYTES, assert_files, mock_http_requests TEST_API_KEY = "test-key" +def fetch_index(args: list[str]) -> int: + return cmd_fetch_index(client=None, args=build_parser().parse_args(["fetch-index"] + args), api_key=None) + + def _govinfo_billstatus(congress: int, btype: str, number: int, *codes: str) -> bytes: """A govinfo BILLSTATUS body whose textVersions carry content/pkg URLs. @@ -973,7 +981,6 @@ def test_valid_content_still_saves(self, tmp_path): def _write_search_corpus(dirpath): """A minimal local BILLSTATUS ZIP with one approps + one non-approps bill.""" - import zipfile def doc(number, title, code): return ( @@ -1037,7 +1044,6 @@ def test_facet_absent_returns_non_appropriations_bills(self, tmp_path, capsys): assert "118-hr-5" in out def test_congress_and_type_filters_narrow_the_index(self, tmp_path, capsys): - import zipfile def doc(congress, btype, number, title): return ( @@ -1158,34 +1164,12 @@ def test_requires_congress(self): with pytest.raises(SystemExit): build_parser().parse_args(["fetch-index"]) - def test_wires_scoped_single_type_download(self, tmp_path, monkeypatch): - import fetch_bills - - calls = {} - - def fake_download(from_congress, to_congress, *, bill_types, destination): - calls.update( - from_congress=from_congress, - to_congress=to_congress, - bill_types=bill_types, - destination=destination, - ) - (destination / "118-hr.zip").write_bytes(b"") # simulate the landed archive - return [destination / "118-hr.zip"] - - monkeypatch.setattr(fetch_bills, "download_archives", fake_download) - args = build_parser().parse_args( - ["fetch-index", "--congress", "118", "--type", "hr", "--billstatus-dir", str(tmp_path)] - ) - from fetch_bills import cmd_fetch_index - - rc = cmd_fetch_index(None, args, None) - assert rc == 0 # every requested archive present after the run - assert calls["from_congress"] == 118 - assert calls["to_congress"] == 118 # single congress: the lightweight slice - assert calls["bill_types"] == ["hr"] - # Resolved against cwd so it points where `search` reads (not script-relative). - assert calls["destination"] == tmp_path.resolve() + @respx.mock + def test_single_congress_and_bill_type_download(self, tmp_path): + mock_http_requests(content=EMPTY_ZIP_BYTES) + rc = fetch_index(["--congress", "118", "--type", "hr", "--billstatus-dir", str(tmp_path)]) + assert rc == 0 + assert_files(tmp_path, {"BILLSTATUS-118-hr.zip"}) def test_type_omitted_fetches_all_types_for_the_congress(self, tmp_path, monkeypatch): import fetch_bills @@ -1244,18 +1228,17 @@ def test_main_propagates_fetch_index_exit_code(self, tmp_path, monkeypatch): fetch_bills.main() assert exc.value.code == 1 + @respx.mock def test_main_routes_and_succeeds(self, tmp_path, monkeypatch): # Happy-path routing: main() dispatches `fetch-index` to the download and exits 0 # when the archive lands. - import fetch_bills - - monkeypatch.setattr(fetch_bills, "download_archives", _fake_download_landing("118-hr.zip")) + mock_http_requests(content=EMPTY_ZIP_BYTES) monkeypatch.setattr( "sys.argv", ["fetch_bills", "fetch-index", "--congress", "118", "--type", "hr", "--billstatus-dir", str(tmp_path)], ) with pytest.raises(SystemExit) as exc: - fetch_bills.main() + fb.main() assert exc.value.code == 0 diff --git a/tests/test_surface_boundary.py b/tests/test_surface_boundary.py index f877decb..a24efae6 100644 --- a/tests/test_surface_boundary.py +++ b/tests/test_surface_boundary.py @@ -176,7 +176,7 @@ def test_the_boundary_scan_actually_looked_at_something(): ) forbidden = _forbidden_names() - for expected in ("fetch_bills", "bill_index", "shared", "web"): + for expected in ("fetch_bills", "shared", "web"): assert expected in forbidden, f"forbidden roster missed {expected!r} -- derivation is broken, not the code" diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 00000000..805a7748 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,66 @@ +"""Shared test helpers.""" + +from __future__ import annotations + +import io +import re +import zipfile +from pathlib import Path + +import respx + +# Validation + + +def assert_files(folder: Path, files: set[str] | list[str]) -> None: + """Assert the folder contains exactly the given filenames.""" + __tracebackhide__ = True + actual = {path.name for path in folder.iterdir()} + expected = set(files) + if actual != expected: + extra = actual - expected + missing = expected - actual + raise AssertionError( + "\n".join( + filter( + None, + [ + f"Unexpected file contents in folder {folder}:", + f"expected: {expected}", + f"actual: {actual}", + f"extra: {extra}" if extra else None, + f"missing: {missing}" if missing else None, + ], + ) + ) + ) + + +# HTTP Mocking +def mock_http_requests( + url: re.Pattern[str] = re.compile(".*"), + status_code: int = 200, + content: bytes = b"", + **kwargs, +) -> None: + """Mock matching GET requests with one response.""" + return respx.get(url).respond(status_code, content=content, **kwargs) + + +# Zip File Mocking +def archive_bytes(members: dict[str, bytes] = {}) -> Path: + """Write a well-formed ZIP named ``{name}.zip`` into source.""" + buf = io.BytesIO() + with zipfile.ZipFile(buf, "w") as zf: + for member, body in members.items(): + zf.writestr(member, body) + return buf.getvalue() + + +EMPTY_ZIP_BYTES = archive_bytes() + + +def write_archive(source: Path, name: str, members: dict[str, bytes] | None = None) -> Path: + """Write a well-formed ZIP named ``{name}.zip`` into source.""" + path = source / name + path.write_bytes(archive_bytes(name, members)) diff --git a/tools/bill_index/__init__.py b/tools/bill_index/__init__.py deleted file mode 100644 index 18ad558c..00000000 --- a/tools/bill_index/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -"""Bill index helpers.""" - -from .bill_index import BillIdentifier, BillIndex, InsertMode, make_bill_id, parse_bill_id - -__all__ = ["BillIndex", "InsertMode", "make_bill_id", "BillIdentifier", "parse_bill_id"] diff --git a/tools/fetch_bill_archives.py b/tools/fetch_bill_archives.py index 13e21996..4ed8c75d 100755 --- a/tools/fetch_bill_archives.py +++ b/tools/fetch_bill_archives.py @@ -2,153 +2,76 @@ """ fetch_bill_archives: -Downloads bulk bill data using GovInfo BILLSTATUS bulk archive ZIP files. -This is a separate API and separate logic from the usual congres.gov API. -It is meant for large volumes of bills. When combined with bill_index, -fetch_bill_archives creates a large index of bill metadata that can be -used for data analysis and testing. +Companion to ``fetch_bill_text_archives.py``. Downloads BILLSTATUS bulk ZIPs and +extracts them into ``bills//_status.xml``. """ from __future__ import annotations +import argparse import re -import shutil import sys import xml.etree.ElementTree as ET import zipfile from datetime import date from pathlib import Path -from time import perf_counter -from typing import Any, Iterator +from typing import Any import httpx -from bill_index import BillIndex, InsertMode, make_bill_id -from shared.bill_types import BILL_TYPES +from shared.bill_index import BillIndex, make_bill_id +from shared.bill_types import resolve_bill_types +from shared.http import download_archives as http_download_archives +from shared.zip import extract_archive -BillMetadata = dict[str, Any] - -# The REPOSITORY root, not this script's directory: `parents[1]` because the fetch -# tooling lives in `tools/` while the working directories it fills are gitignored at the -# root (`/bills`, `/bills_bulk_text` — anchored, see .gitignore). Resolving them beside -# the script instead would put hundreds of MB of downloads outside those rules, which is -# the silent-`git add` failure #308 exists to prevent (#367). PROJECT_DIR = Path(__file__).resolve().parents[1] - DEFAULT_BILLS_DIR = PROJECT_DIR / "bills" +DEFAULT_ZIP_DIR = PROJECT_DIR / "bills_bulk_status" +BILLSTATUS_ZIP_FORMAT = "BILLSTATUS-{congress}-{bill_type}.zip" -GOVINFO_BILLSTATUS_ZIP_URL = ( - "https://www.govinfo.gov/bulkdata/BILLSTATUS/{congress}/{bill_type}/BILLSTATUS-{congress}-{bill_type}.zip" +GOVINFO_BASE_URL = "https://www.govinfo.gov/bulkdata/" +GOVINFO_BILLSTATUS_ZIP_URL_FORMAT = ( + GOVINFO_BASE_URL + "BILLSTATUS/{congress}/{bill_type}/BILLSTATUS-{congress}-{bill_type}.zip" ) - -# Ceiling on what one archive may expand to on disk (#279). `zipfile.extractall` -# applies no bound of its own, so a crafted archive -- a few MB of highly repetitive -# data that inflates to terabytes -- would fill the disk before anything noticed. -# -# Calibrated against the real bulk data rather than guessed: on 2026-07-21 the largest -# BILLSTATUS archive is 118-hr, 34 MB compressed expanding to 162 MiB across 10,564 -# members, a 5.0x ratio; every member is XML. 2 GiB leaves better than an order of -# magnitude of headroom for corpus growth while still refusing anything that could -# plausibly exhaust a disk. Raise it deliberately if real archives ever approach it; -# the failure mode of a too-low ceiling is a loud refusal, not a silent truncation. -MAX_UNCOMPRESSED_BYTES = 2 * 1024**3 - -# Companion ceiling on member *count* (#306). The byte ceiling above sums declared -# uncompressed sizes, but empty members declare zero bytes: 200,000 empty files weigh -# nothing against MAX_UNCOMPRESSED_BYTES yet still consume 200,000 inodes and directory -# entries, exhausting the filesystem the byte ceiling was meant to protect. Inode -# exhaustion degrades worse than a full disk -- it affects the whole machine and is -# harder to diagnose, since free space still reads as available. -# -# Calibrated the same way as the byte ceiling: the largest real BILLSTATUS archive is -# 118-hr at 10,564 members (the figure #300 measured), so 100,000 keeps better than a -# 9x margin for corpus growth while refusing the 200,000-member archive #306 demonstrated -# by a factor of two. As with the byte ceiling, a too-low bound fails loud (a refusal), -# never silent -- raise it deliberately if real archives ever approach it. -MAX_MEMBER_COUNT = 100_000 - -_POPULAR_TITLE_RE = re.compile(r"^popular\s+titles?\b", re.IGNORECASE) -_BILLSTATUS_XML_GLOB = "BILLSTATUS*.xml" -_BILLSTATUS_XML_NAME_RE = re.compile(r"^BILLSTATUS-(\d+)([a-z]+)(\d+)\.xml$", re.IGNORECASE) -LOG_PERFORMANCE = False - -_LEGACY_COLUMN_RENAMES = { - "action_count": "actionCount", - "version_count": "versionCount", - "budget_estimate_count": "budgetEstimateCount", - "amendment_count": "amendmentCount", - "related_bills_count": "relatedBillsCount", -} - - -# STEP 1: Download zip archives -# The fastest way to get massive amounts of bill metadata is to download complete zip archives from govinfo bulk data. -# Bill archives are stored per congress and bill type. Each zip file has per-bill metadata for up to thousands of bills. -def archive_url(congress: int, bill_type: str) -> str: - """Build direct download URL for one BILLSTATUS archive ZIP.""" - return GOVINFO_BILLSTATUS_ZIP_URL.format(congress=congress, bill_type=bill_type) - - -def resolve_destination(destination: Path | str | None = None) -> Path: - """Resolve a relative destination against the repository root (see PROJECT_DIR).""" - path = Path(destination or DEFAULT_BILLS_DIR) - if not path.is_absolute(): - path = PROJECT_DIR / path - return path - - -def archive_destination(destination: Path, congress: int, bill_type: str) -> Path: - """Build the output path for one BILLSTATUS archive ZIP.""" - return destination / f"{congress}-{bill_type}.zip" +GOVINFO_BILL_FILENAME_RE = re.compile( + r"^BILLSTATUS-(\d+)([a-z]+)(\d+)\.xml$", + re.IGNORECASE, +) +GOVINFO_BILLSTATUS_FILENAME_FORMAT = "BILLSTATUS-{congress}{bill_type}{number}.xml" -def archive_error_path(destination: Path, congress: int, bill_type: str) -> Path: - """Build the error marker path for one failed archive download.""" - return destination / f"{congress}-{bill_type}.error" +def parse_billstatus_filename(filename: str) -> tuple[int, str, int]: # (congress, bill_type, number) + """``BILLSTATUS-119hr1.xml`` → ``(119, "hr", 1)``.""" + match = GOVINFO_BILL_FILENAME_RE.match(Path(filename).name) + if not match: + return (0, "", 0) + congress, bill_type, number = match.groups() + return int(congress), bill_type, int(number) -def _print_download_progress(downloaded: int, total: int) -> None: - """Print a single-line download progress update to stderr.""" - if total: - pct = downloaded * 100 // total - mb_done = downloaded / (1024 * 1024) - mb_total = total / (1024 * 1024) - print(f"\r {mb_done:.1f}/{mb_total:.1f} MB ({pct}%)", end="", file=sys.stderr, flush=True) - else: - mb_done = downloaded / (1024 * 1024) - print(f"\r {mb_done:.1f} MB downloaded", end="", file=sys.stderr, flush=True) +def archive_destination(destination: Path, congress: int, bill_type: str) -> Path: + """Return the local path for one BILLSTATUS archive.""" + return destination / BILLSTATUS_ZIP_FORMAT.format(congress=congress, bill_type=bill_type) -def archive_temp_path(dest: Path) -> Path: - """Build temporary path used while downloading one archive.""" - return dest.with_suffix(dest.suffix + ".part") +def billstatus_zip_url(congress: int, bill_type: str) -> str: + return GOVINFO_BILLSTATUS_ZIP_URL_FORMAT.format(congress=congress, bill_type=bill_type) -def _verify_archive_complete(path: Path) -> None: - """Raise unless path is a readable ZIP archive. +def billstatus_zip_filename(congress: int, bill_type: str) -> str: + return BILLSTATUS_ZIP_FORMAT.format(congress=congress, bill_type=bill_type) - The content-length check is the completeness signal only when the server sends - that header; a chunked response legitimately omits it, and then a truncated body - is indistinguishable from a whole one by byte count alone (#63). The archive's own - end-of-central-directory record is the fallback signal: it is written last, so a - short read loses it and the file no longer opens. This is the same operation - extract_archive performs downstream -- doing it before committing turns a silently - cached partial archive into a failed download that the next run retries. - Emptiness is deliberately not checked: a zero-member ZIP is structurally valid, - and truncation always destroys the end-of-central-directory record, so a short - read can only ever produce "does not open", never "opens with zero members". - """ - try: - with zipfile.ZipFile(path): - pass - except (zipfile.BadZipFile, OSError) as exc: - raise httpx.HTTPError(f"Incomplete download: {path.name} is not a readable ZIP archive ({exc})") from exc +def billstatus_filename(congress: int, bill_type: str, number: int) -> str: + return GOVINFO_BILLSTATUS_FILENAME_FORMAT.format(congress=congress, bill_type=bill_type, number=number) -def _progress_prefix(index: int, total: int) -> str: - """Build a ``current/total:`` progress prefix for batch status lines.""" - return f"{index}/{total}:" +def enumerate_congresses(from_congress: int, to_congress: int) -> list[int]: + return list( + range(from_congress, to_congress + 1) + if from_congress <= to_congress + else range(from_congress, to_congress - 1, -1) + ) def enumerate_tasks( @@ -157,426 +80,160 @@ def enumerate_tasks( *, bill_types: list[str] | None = None, ) -> list[tuple[int, str]]: - """Return newest-first (congress, bill_type) tasks for a validated selection.""" - selected_types = _validate_archive_params(from_congress, to_congress, bill_types) - congresses = reversed(range(from_congress, to_congress + 1)) - return [(congress, bill_type) for congress in congresses for bill_type in selected_types] - - -# STEP 1: Download archives -# The bulk data API has zip files per congress and bill type. Each ZIP file has per-bill metadata. -def download_archive_zip(client: httpx.Client, url: str, dest: Path) -> None: - """Stream one archive ZIP to disk atomically, printing progress when possible.""" - temp_path = archive_temp_path(dest) - if temp_path.exists(): - temp_path.unlink() - - try: - with client.stream("GET", url, follow_redirects=True, timeout=300) as response: - response.raise_for_status() - total = int(response.headers.get("content-length", 0) or 0) - downloaded = 0 - with temp_path.open("wb") as fh: - for chunk in response.iter_bytes(chunk_size=256 * 1024): - if not chunk: - continue - fh.write(chunk) - downloaded += len(chunk) - _print_download_progress(downloaded, total) - print(file=sys.stderr) - if total and downloaded != total: - raise httpx.HTTPError(f"Incomplete download: got {downloaded} of {total} bytes") - _verify_archive_complete(temp_path) - temp_path.replace(dest) - except Exception: - if temp_path.exists(): - temp_path.unlink() - raise + """Return the BILLSTATUS archive scopes for a congress range.""" + return [ + (congress, bill_type) + for congress in enumerate_congresses(from_congress, to_congress) + for bill_type in resolve_bill_types(bill_types) + ] def download_archives( from_congress: int, to_congress: int, + bill_types: list[str], + destination: Path, *, - bill_types: list[str] | None = None, - destination: Path | str | None = None, + overwrite_existing: bool = False, ) -> list[Path]: - """Download BILLSTATUS archive ZIPs for congress/type combinations.""" - destination = resolve_destination(destination) - destination.mkdir(parents=True, exist_ok=True) - downloaded: list[Path] = [] + """Download BILLSTATUS ZIPs for each (congress, type); skip existing unless overwriting.""" tasks = enumerate_tasks(from_congress, to_congress, bill_types=bill_types) - total = len(tasks) - + urls = [billstatus_zip_url(congress, bill_type) for congress, bill_type in tasks] with httpx.Client(timeout=300) as client: - for index, (congress, bill_type) in enumerate(tasks, start=1): - prefix = _progress_prefix(index, total) - dest = archive_destination(destination, congress, bill_type) - error_path = archive_error_path(destination, congress, bill_type) - - if dest.exists(): - print(f"{prefix} Skipping existing archive: {dest.name}", file=sys.stderr) - # The archive beside it disproves the marker, so clear it here too and - # not only on the download path (#259). Leaving it would let the two - # states contradict each other for as long as the cache survives. - error_path.unlink(missing_ok=True) - continue - - url = archive_url(congress, bill_type) - print(f"{prefix} Downloading {dest.name}", file=sys.stderr) - print(f" {url}", file=sys.stderr) - try: - download_archive_zip(client, url, dest) - except Exception as exc: - error_path.write_text(str(exc), encoding="utf-8") - print(f"{prefix} Failed {dest.name}: wrote {error_path.name}", file=sys.stderr) - continue - - if error_path.exists(): - error_path.unlink() - downloaded.append(dest) - print(f"{prefix} Saved: {dest.name}", file=sys.stderr) - - return downloaded - - -# Step 2: Extract archives -# Once we have zip files per congress and bill type, we need to extract the per-bill metadata from the archives. -def archive_extract_dir(source: Path, archive: Path) -> Path: - """Build extraction directory for one archive (same stem as the zip).""" - return source / archive.stem - - -def extract_archive(archive: Path, dest_dir: Path) -> None: - """Extract one archive ZIP into dest_dir, refusing an oversized expansion. - - Two ceilings guard extraction, both read from the central directory before a - single byte is written -- a ceiling enforced during extraction has already spent - the disk it protects. The byte ceiling bounds total uncompressed size; the member - ceiling bounds file *count*, because empty members declare zero bytes and so slip - the byte ceiling entirely while still consuming inodes (#306). Member *paths* are - untrusted too, but `zipfile.extractall` already sanitizes traversal and absolute - paths (see test_members_cannot_escape_the_destination_directory), so - re-implementing the walk to add filtering would reintroduce that escape to solve a - problem the stdlib has handled. - """ - with zipfile.ZipFile(archive) as zf: - infos = zf.infolist() - member_count = len(infos) - if member_count > MAX_MEMBER_COUNT: - raise ValueError( - f"{archive.name} holds {member_count} members, over the {MAX_MEMBER_COUNT} ceiling; refusing to extract" - ) - declared = sum(info.file_size for info in infos) - if declared > MAX_UNCOMPRESSED_BYTES: - raise ValueError( - f"{archive.name} declares {declared} uncompressed bytes, over the " - f"{MAX_UNCOMPRESSED_BYTES} ceiling; refusing to extract" - ) - dest_dir.mkdir(parents=True, exist_ok=True) - zf.extractall(dest_dir) - - -def extract_archives(source: Path | str | None = None) -> list[Path]: - """Extract all ZIP archives in source, skipping existing folders.""" - source = resolve_destination(source) - if not source.is_dir(): - raise ValueError(f"Source folder does not exist: {source}") - - extracted: list[Path] = [] - archives = sorted(source.glob("*.zip")) - total = len(archives) - - for index, archive in enumerate(archives, start=1): - prefix = _progress_prefix(index, total) - dest_dir = archive_extract_dir(source, archive) - if dest_dir.exists(): - print(f"{prefix} Skipping existing folder: {dest_dir.name}", file=sys.stderr) - continue - - print(f"{prefix} Extracting {archive.name}", file=sys.stderr) - try: - extract_archive(archive, dest_dir) - except Exception as exc: - if dest_dir.exists(): - shutil.rmtree(dest_dir) - print(f"{prefix} Failed {archive.name}: {exc}", file=sys.stderr) - continue - - extracted.append(dest_dir) - - return extracted - - -def iter_billstatus_files( - from_congress: int, - to_congress: int, - *, - bill_types: list[str] | None = None, - destination: Path | str | None = None, -) -> Iterator[Path]: - """Yield BILLSTATUS XML files for archive folders matching the congress/type selection.""" - for congress, bill_type in enumerate_tasks( - from_congress, - to_congress, - bill_types=bill_types, - ): - yield from enumerate_files(congress, bill_type, destination=destination) - - -def enumerate_files( - congress: int, - bill_type: str, - *, - destination: Path | str | None = None, -) -> list[Path]: - """Return BILLSTATUS XML files for one extracted archive folder.""" - destination = resolve_destination(destination) - archive_dir = archive_extract_dir( - destination, - archive_destination(destination, congress, bill_type), - ) - if not archive_dir.is_dir(): - return [] - return sorted(archive_dir.glob(_BILLSTATUS_XML_GLOB)) - - -def _bill_id_from_xml_path(xml_path: Path | str) -> str | None: - """Derive bill id from BILLSTATUS xml file name without opening the file.""" - match = _BILLSTATUS_XML_NAME_RE.match(Path(xml_path).name) - if not match: - return None - congress, bill_type, number = match.groups() - return make_bill_id(congress, bill_type.lower(), number) - - -def _xml_path_from_bill_id(bill_id: str) -> Path: - """Derive BILLSTATUS XML file path from bill id.""" - congress, bill_type, number = bill_id.split("-") - return Path(f"BILLSTATUS-{congress}{bill_type}{number}.xml") - - -def _pick_bill_title(bill: ET.Element) -> str: - """Prefer a popular title; otherwise use the first listed title.""" - titles = bill.find("titles") - if titles is not None: - for item in titles.findall("item"): - title_type = item.findtext("titleType", "") - if _POPULAR_TITLE_RE.match(title_type): - title = item.findtext("title", "").strip() - if title: - return title - - for item in titles.findall("item"): - title = item.findtext("title", "").strip() - if title: - return title - - return bill.findtext("title", "").strip() - - -def _bill_number(bill: ET.Element) -> str: - """Read bill number from modern or legacy BILLSTATUS XML.""" - return (bill.findtext("number") or bill.findtext("billNumber") or "").strip() - - -def _bill_type_slug(bill: ET.Element) -> str: - """Read bill type slug from modern or legacy BILLSTATUS XML.""" - return (bill.findtext("type") or bill.findtext("billType") or "").strip().lower() - - -def _committee_count(bill: ET.Element) -> int: - """Count committees across modern and legacy BILLSTATUS XML layouts.""" - items = bill.findall("committees/item") - if not items: - items = bill.findall("committees/billCommittees/item") - return len(items) - - -def _first_summary_length(bill: ET.Element) -> int: - """Return character length of the first CRS summary text, if present.""" - summaries = bill.find("summaries") - if summaries is None: - return 0 - - for tag_path in ("summary", "billSummaries/item", "item"): - first_summary = summaries.find(tag_path) - if first_summary is not None: - return len((first_summary.findtext("text", "") or "").strip()) - - return 0 - - -def _days_active(introduced_date: str, last_action_date: str) -> int | None: - """Return days between introduction and last action, if both dates are present.""" - if not introduced_date or not last_action_date: - return None - start = date.fromisoformat(introduced_date) - end = date.fromisoformat(last_action_date) - return (end - start).days - + return http_download_archives( + client, + urls, + destination, + url_to_path=lambda url, index: archive_destination(Path(), *tasks[index]), + skip_existing=not overwrite_existing, + ) -def extract_bill_metadata_from_archive_xml(source: Path | str) -> BillMetadata: - """ - Convert one GovInfo BILLSTATUS XML file into succinct bill metadata. - XML data is from responses to requests of the form: - https://www.govinfo.gov/bulkdata/BILLSTATUS/119/hr/BILLSTATUS-119hr123.xml - """ - xml_path = Path(source) - bill = ET.parse(xml_path).getroot().find("bill") +def extract_bill_metadata(xml_content: str | bytes, bill_id: str) -> dict[str, Any]: + """Pull a short status summary from one BILLSTATUS XML. ``bill_id`` comes from the filename.""" + if isinstance(xml_content, bytes): + xml_content = xml_content.decode("utf-8", errors="replace") + bill = ET.fromstring(xml_content).find("bill") if bill is None: - raise ValueError(f"No element found in {xml_path}") - - congress = bill.findtext("congress", "").strip() - number = _bill_number(bill) - bill_type = _bill_type_slug(bill) - if not congress or not number or not bill_type: - raise ValueError(f"Missing congress, type, or number in {xml_path}") + raise ValueError(f"No element in {bill_id}") - introduced_date = bill.findtext("introducedDate", "").strip() - last_action_date = bill.findtext("latestAction/actionDate", "").strip() + introduced = bill.findtext("introducedDate", "").strip() + last_action = bill.findtext("latestAction/actionDate", "").strip() + days_active = None + if introduced and last_action: + days_active = (date.fromisoformat(last_action) - date.fromisoformat(introduced)).days + committees = bill.findall("committees/item") or bill.findall("committees/billCommittees/item") return { - "id": make_bill_id(congress, bill_type, number), - "title": _pick_bill_title(bill), - "introducedDate": introduced_date, - "lastActionDate": last_action_date, - "daysActive": _days_active(introduced_date, last_action_date), + "id": bill_id, + "title": (bill.findtext("title") or "").strip(), + "introducedDate": introduced, + "lastActionDate": last_action, + "daysActive": days_active, "status": bill.findtext("latestAction/text", "").strip(), "policyArea": bill.findtext("policyArea/name", "").strip(), - "historySize": xml_path.stat().st_size, - "summaryLength": _first_summary_length(bill), + "historySize": len(xml_content), "actionCount": len(bill.findall("actions/item")), "versionCount": len(bill.findall("textVersions/item")), - "budgetEstimateCount": len(bill.findall("cboCostEstimates/item")), "amendmentCount": len(bill.findall("amendments/amendment")), "relatedBillsCount": len(bill.findall("relatedBills/item")), - "committeeCount": _committee_count(bill), - "sponsorCount": len(bill.findall("sponsors/item")), + "committeeCount": len(committees), } -def _validate_archive_params( - from_congress: int, - to_congress: int, - bill_types: list[str] | None, -) -> list[str]: - """Validate congress range and bill types; return selected type slugs.""" - if from_congress > to_congress: - raise ValueError(f"from_congress ({from_congress}) must be <= to_congress ({to_congress})") - - selected_types = bill_types or list(BILL_TYPES) - unknown = [bill_type for bill_type in selected_types if bill_type not in BILL_TYPES] - if unknown: - raise ValueError(f"Unknown bill types: {unknown}") - return selected_types - - -def parse_bill_archives( +def convert_archives( + zip_dir: Path, + out_dir: Path, + *, from_congress: int, to_congress: int, - *, - bill_types: list[str] | None = None, - destination: Path | str | None = None, - index: BillIndex | None = None, - mode: InsertMode = "skip", -): - """Parse BILLSTATUS XML for archive folders matching the congress/type selection.""" - destination = resolve_destination(destination) + bill_types: list[str], + overwrite_existing: bool = False, + bill_index_path: Path | None = None, +) -> None: + """Extract BILLSTATUS members via ``extract_archive`` into ``/_status.xml``.""" tasks = enumerate_tasks(from_congress, to_congress, bill_types=bill_types) - task_count = len(tasks) - index = index or BillIndex(DEFAULT_BILLS_DIR / "bills.csv") - index.rename_columns(_LEGACY_COLUMN_RENAMES) - for task_index, (congress, bill_type) in enumerate(tasks, start=1): - prefix = _progress_prefix(task_index, task_count) - - enum_start = perf_counter() - bill_xml_paths = enumerate_files(congress, bill_type, destination=destination) - enumerate_secs = perf_counter() - enum_start - bill_ids = [_bill_id_from_xml_path(xml_path) for xml_path in bill_xml_paths] - bill_paths_by_id = {bill_id: xml_path for xml_path, bill_id in zip(bill_xml_paths, bill_ids)} - new_bill_ids, existing_bill_ids = index.find_new_and_existing_bill_ids(bill_ids) - parse_bill_ids = new_bill_ids if mode == "skip" else bill_ids - parse_bill_paths = [bill_paths_by_id[bill_id] for bill_id in parse_bill_ids] - - extract_start = perf_counter() - records = [extract_bill_metadata_from_archive_xml(xml_path) for xml_path in parse_bill_paths] - extract_secs = perf_counter() - extract_start - - merge_start = perf_counter() - index.add_bills(records, mode=mode) - merge_secs = perf_counter() - merge_start - - status_parts = [] - if existing_bill_ids: - status_parts.append(f"found {len(existing_bill_ids)} existing bills") - if new_bill_ids: - status_parts.append(f"added {len(new_bill_ids)} new bills") - updated_count = len(existing_bill_ids) if mode != "skip" else 0 - if updated_count: - status_parts.append(f"updated {updated_count} bills") - - print( - f"{prefix} {congress}-{bill_type} - {', '.join(status_parts)}", - file=sys.stderr, - ) + zip_paths = [archive_destination(zip_dir, congress, bill_type) for congress, bill_type in tasks] + bill_index = BillIndex(csv_path=bill_index_path) if bill_index_path else None + records: list[dict[str, Any]] = [] + + def handle_file(filename: str, _i: int, _zf: zipfile.ZipFile) -> str | None: + congress, bill_type, number = parse_billstatus_filename(filename) + bill_id = f"{congress}-{bill_type}-{number}" + return f"{bill_id}/{bill_id}_status.xml" + + def handle_content(content: bytes, filename: str, _i: int, _zf: zipfile.ZipFile) -> bytes: + if bill_index is not None: + congress, bill_type, number = parse_billstatus_filename(filename) + bill_id = make_bill_id(congress, bill_type, number) + try: + records.append(extract_bill_metadata(content, bill_id)) + except Exception as exc: + print(f" error extracting metadata from {filename}: {exc}", file=sys.stderr) + records.append({"id": bill_id, "bill_status_error": str(exc)}) + return content - if LOG_PERFORMANCE: - print( - ( - f"{prefix} {congress}-{bill_type} performance - " - f"enumerating files: {enumerate_secs:.3f}s, " - f"extracting metadata: {extract_secs:.3f}s, " - f"merging index: {merge_secs:.3f}s" - ), - file=sys.stderr, + for i, path in enumerate(zip_paths): + print(f" {i + 1}/{len(zip_paths)}: extracting {path.name}...", file=sys.stderr) + try: + extract_archive( + path, + out_dir=out_dir, + files=GOVINFO_BILL_FILENAME_RE, + overwrite_existing=overwrite_existing, + file_handler=handle_file, + file_content_handler=handle_content, ) - - if index is not None: - print( - (f"Done parsing: bill index saved at {index.csv_path.resolve()} with {len(index.bills)} records"), - file=sys.stderr, - ) + except Exception as exc: + if not path.exists(): + print(f" {path.name} not found, skipping", file=sys.stderr) + else: + print(f" error extracting {path.name}: {exc}", file=sys.stderr) + + if bill_index is not None and records: + bill_index.add_bills(records, mode="merge") + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--from-congress", type=int, default=118) + p.add_argument("--to-congress", type=int, default=119) + p.add_argument( + "--types", + nargs="+", + default=["all"], + help="Bill types to download. Use 'all' to include every key from shared/BILL_TYPES.", + ) + p.add_argument("--zip-dir", type=Path, default=DEFAULT_ZIP_DIR) + p.add_argument("--out-dir", type=Path, default=DEFAULT_BILLS_DIR) + p.add_argument("--bill-index-file", type=Path, help="Optional CSV to merge per-bill metadata") + p.add_argument( + "--overwrite-existing", + action="store_true", + help="Re-download ZIPs and overwrite extracted status files (default: skip existing)", + ) + return p -def fetch_bill_archives( - from_congress: int, - to_congress: int, - *, - bill_types: list[str] | None = None, - destination: Path | str | None = None, - index: BillIndex | None = None, - mode: InsertMode = "merge", -) -> list[BillMetadata]: - """Download, extract, and index GovInfo BILLSTATUS bulk archives. - - Each phase skips work that is already done: existing ZIPs, extracted folders, - and bill ids already present in the index (when merging). - """ - destination = resolve_destination(destination) - - print("Phase 1/3: Download archives", file=sys.stderr) +def main() -> None: + args = build_parser().parse_args() download_archives( - from_congress, - to_congress, - bill_types=bill_types, - destination=destination, + args.from_congress, + args.to_congress, + args.types, + args.zip_dir, + overwrite_existing=args.overwrite_existing, ) - - print("Phase 2/3: Extract archives", file=sys.stderr) - extract_archives(destination) - - print("Phase 3/3: Parse metadata into index", file=sys.stderr) - return parse_bill_archives( - from_congress, - to_congress, - bill_types=bill_types, - destination=destination, - index=index, - mode=mode, + convert_archives( + args.zip_dir, + args.out_dir, + from_congress=args.from_congress, + to_congress=args.to_congress, + bill_types=args.types, + overwrite_existing=args.overwrite_existing, + bill_index_path=args.bill_index_file, ) if __name__ == "__main__": - fetch_bill_archives(112, 119, index=BillIndex(DEFAULT_BILLS_DIR / "bills.csv")) + main() diff --git a/tools/fetch_bills.py b/tools/fetch_bills.py index 480f6ce0..26891f7b 100755 --- a/tools/fetch_bills.py +++ b/tools/fetch_bills.py @@ -19,8 +19,8 @@ from dotenv import load_dotenv import fetch_govinfo as gi -from bill_index import BillIdentifier, BillIndex, parse_bill_id from fetch_bill_archives import archive_destination, download_archives, enumerate_tasks +from shared.bill_index import BillIdentifier, BillIndex, parse_bill_id from shared.bill_types import BILL_TYPES from shared.http import api_get, request_with_retry diff --git a/tools/bill_index/bill_index.py b/tools/shared/bill_index.py similarity index 100% rename from tools/bill_index/bill_index.py rename to tools/shared/bill_index.py diff --git a/tools/shared/bill_types.py b/tools/shared/bill_types.py index c3df044d..50dce062 100644 --- a/tools/shared/bill_types.py +++ b/tools/shared/bill_types.py @@ -11,3 +11,8 @@ "hconres": ("H.Con.Res.", "house-concurrent-resolution"), "sconres": ("S.Con.Res.", "senate-concurrent-resolution"), } + + +def resolve_bill_types(bill_types: list[str] | None = None) -> list[str]: + "Allow for 'all' keyword to include all bill types. Default to all types if no type are specified." + return list(BILL_TYPES) if bill_types is None or "all" in bill_types else bill_types diff --git a/tools/shared/http.py b/tools/shared/http.py index ab354036..f9f7f821 100644 --- a/tools/shared/http.py +++ b/tools/shared/http.py @@ -4,9 +4,13 @@ import sys import time +from pathlib import Path +from typing import Callable, Iterable import httpx +from shared.zip import verify_archive_complete + BASE_URL = "https://api.congress.gov/v3" LOG_API_REQUESTS = True @@ -67,3 +71,148 @@ def api_get( request_params["api_key"] = api_key resp = request_with_retry(client, url, request_params) return resp.json() + + +def _print_download_progress(downloaded: int, total: int) -> None: + """Print a single-line download progress update to stderr.""" + mb_done = downloaded / (1024 * 1024) + if total: + pct = downloaded * 100 // total + mb_total = total / (1024 * 1024) + print(f"\r {mb_done:.1f}/{mb_total:.1f} MB ({pct}%)", end="", file=sys.stderr, flush=True) + else: + print(f"\r {mb_done:.1f} MB", end="", file=sys.stderr, flush=True) + + +def path_for_error(path: Path) -> Path: + """Return the error-marker path associated with ``path``.""" + return path.with_suffix(path.suffix + ".error") + + +def write_error(error: Exception, path: Path) -> Path: + """Write an error beside ``path`` and return the marker path.""" + error_path = path_for_error(Path(path)) + error_path.parent.mkdir(parents=True, exist_ok=True) + error_path.write_text(str(error), encoding="utf-8") + return error_path + + +def download_temp_path(destination: Path) -> Path: + return destination.with_suffix(destination.suffix + ".part") + + +def cached_file_download( + client: httpx.Client, + url: str, + destination: Path, + *, + skip_existing: bool = True, + verify: Callable[[Path], None] | None = None, +) -> bool: + """Stream ``url`` to ``destination`` atomically, skipping it if already present. + Returns True if downloaded and False if skipped. Raises when downloading fails. + + Creates a temporary '.part' file while downloading to ensure atomic downloads. + """ + error_path = path_for_error(destination) + if skip_existing and destination.exists(): + error_path.unlink(missing_ok=True) + return False + + destination.parent.mkdir(parents=True, exist_ok=True) + temp_path = download_temp_path(destination) + if temp_path.exists(): + temp_path.unlink() + + try: + with client.stream("GET", url, follow_redirects=True, timeout=300) as response: + response.raise_for_status() + total = int(response.headers.get("content-length", 0) or 0) + downloaded = 0 + with temp_path.open("wb") as fh: + for chunk in response.iter_bytes(chunk_size=256 * 1024): + if not chunk: + continue + fh.write(chunk) + downloaded += len(chunk) + _print_download_progress(downloaded, total) + print(file=sys.stderr) + if total and downloaded != total: + raise httpx.HTTPError(f"Incomplete download: got {downloaded} of {total} bytes") + + # Verify before committing the temp file into place. This preserves the + # atomicity invariant: a failed download must not leave a destination + # file behind. + if verify is not None: + verify(temp_path) + temp_path.replace(destination) + error_path.unlink(missing_ok=True) + return True + except Exception as error: + if temp_path.exists(): + temp_path.unlink() + write_error(error, destination) + raise + + +def download_zip( + client: httpx.Client, + url: str, + destination: Path, + *, + skip_existing: bool = True, +) -> bool: + """Download a ZIP atomically, verifying completeness before committing.""" + return cached_file_download( + client, + url, + destination, + skip_existing=skip_existing, + verify=verify_archive_complete, + ) + + +def download_archives( + client: httpx.Client, + urls: Iterable[str], + destination: Path, + *, + url_to_path: Callable[[str, int], Path], + skip_existing: bool = True, +) -> list[Path]: + """ + Download many archive URLs into ``destination``, continuing past failures. + """ + destination = Path(destination) + destination.mkdir(parents=True, exist_ok=True) + url_list = list(urls) + saved: list[Path] = [] + + print(f"Downloading {len(url_list)} zip files...", file=sys.stderr) + for i, url in enumerate(url_list): + mapped = Path(url_to_path(url, i)) + dest = mapped if mapped.is_absolute() else destination / mapped + prefix = f"{i + 1}/{len(url_list)}:" + try: + downloaded = download_zip( + client, + url, + dest, + skip_existing=skip_existing, + ) + except httpx.HTTPStatusError as exc: + if exc.response.status_code == 404: + print(f"{prefix} no zip (404) for {dest.name}", file=sys.stderr) + else: + print(f"{prefix} FAILED {dest.name}: {exc}", file=sys.stderr) + continue + except Exception as exc: + print(f"{prefix} FAILED {dest.name}: {exc}", file=sys.stderr) + continue + + if downloaded: + saved.append(dest) + print(f"{prefix} saved {dest.name} ({url})", file=sys.stderr) + else: + print(f"{prefix} skip existing {dest.name}", file=sys.stderr) + return saved diff --git a/tools/shared/zip.py b/tools/shared/zip.py new file mode 100644 index 00000000..9ba56349 --- /dev/null +++ b/tools/shared/zip.py @@ -0,0 +1,142 @@ +"""Shared ZIP helpers.""" + +from __future__ import annotations + +import fnmatch +import re +import zipfile +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Iterator, NamedTuple + +import httpx + + +@dataclass +class ArchiveFile: + zip_path: str | Path + file_path: str | Path + + def __str__(self) -> str: + return f"{self.zip_path}/{self.file_path}" + + def __repr__(self) -> str: + return f"ArchiveFile({self})" + + +class ExtractArchiveDetails(NamedTuple): + files_extracted: list[Path] + files_skipped: list[Path] + errors: dict[Path, Exception] + + +def _ensure_within_destination(out_dir: Path, dest: Path) -> None: + """Raise ValueError if ``dest`` resolves outside ``out_dir``.""" + out_resolved = out_dir.resolve() + dest_resolved = dest.resolve() + if not dest_resolved.is_relative_to(out_resolved): + raise ValueError(f"Zip member escapes destination directory: {dest}") + + +def verify_archive_complete(path: Path) -> None: + """Raise unless ``path`` is a readable ZIP archive. + + A content-length check is the completeness signal only when the server sends + that header; a chunked response legitimately omits it, and then a truncated + body is indistinguishable from a whole one by byte count alone. The archive's + own end-of-central-directory record is the fallback signal: it is written + last, so a short read loses it and the file no longer opens. + + Emptiness is deliberately not checked: a zero-member ZIP is structurally + valid, and truncation always destroys the end-of-central-directory record, so + a short read can only ever produce "does not open", never "opens with zero + members". + """ + try: + with zipfile.ZipFile(path): + pass + except (zipfile.BadZipFile, OSError) as exc: + raise httpx.HTTPError(f"Incomplete download: {path.name} is not a readable ZIP archive ({exc})") from exc + + +def iterate_archive(path: Path, pattern: str | re.Pattern[str] = "*") -> Iterator[tuple[str, zipfile.ZipFile]]: + r"""Yield ``(path, zip_handle)`` for each archive file or directory matching ``pattern``. + + A string matching pattern uses a shell-style glob with simplified regex semantics. + A pattern of type re.Pattern uses full regex matching. + For example, pattern = "*.xml" is equivalent to pattern = re.compile(r"\.xml") + and pattern = re.compile(r"^.*\.xml$") + """ + with zipfile.ZipFile(path) as zf: + + def _matches(name: str) -> bool: + if isinstance(pattern, re.Pattern): + return pattern.match(Path(name).name) is not None + return fnmatch.fnmatch(name, pattern) + + files = [name for name in zf.namelist() if _matches(name)] + yield from [(name, zf) for name in files] + + +def extract_archive( + archive_path: Path | str, + *, + out_dir: Path | str, + files: str | re.Pattern[str] = "*", + overwrite_existing: bool = False, + file_handler: Callable[[str, int, zipfile.ZipFile], str | Path | None] = lambda filename, index, zf: filename, + file_content_handler: Callable[ + [bytes, str, int, zipfile.ZipFile], bytes | None + ] = lambda data, filename, index, zf: data, +) -> tuple[int, ExtractArchiveDetails]: + """Extract matching ZIP members into ``out_dir``. + + Args: + archive_path: ZIP file to read. + out_dir: Destination root for extracted files. + files: Glob string or compiled regex selecting archive members. + overwrite_existing: When false, skip members whose destination already exists. + file_handler: Maps archive member path to a path relative to ``out_dir`` + to allow the file structure to be reordered. Return ``None`` to skip members. + file_content_handler: Transforms or analyzes file contents before writing. + + Returns: + ``(count, details)`` where ``count`` is how many files were written and + ``details`` has ``files_extracted``, ``files_skipped``, and ``errors``. + """ + out_dir = Path(out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + files_extracted: list[Path] = [] + files_skipped: list[Path] = [] + errors: dict[Path, Exception] = {} + + for index, (name, zf) in enumerate(iterate_archive(Path(archive_path), files)): + member_path = Path(name) + try: + if name.endswith("/"): + files_skipped.append(member_path) + continue + dest_rel = file_handler(name, index, zf) if file_handler else None + if dest_rel is None: + files_skipped.append(member_path) + continue + dest = out_dir / dest_rel + # Check against malicious contents that attempt to escape the destination directory. + # To do: check against other malicious contents like symlinks leading to a bad destination? + _ensure_within_destination(out_dir, dest) + if dest.exists() and not overwrite_existing: + files_skipped.append(dest) + continue + data = file_content_handler(zf.read(name), name, index, zf) if file_content_handler else None + if data is None: + files_skipped.append(dest) + continue + dest.parent.mkdir(parents=True, exist_ok=True) + _ensure_within_destination(out_dir, dest) + dest.write_bytes(data) + files_extracted.append(dest) + except Exception as exc: + errors[member_path] = exc + + return len(files_extracted), ExtractArchiveDetails(files_extracted, files_skipped, errors)