From 61ac99095ae024b6f8ca42171dddf24e964bf645 Mon Sep 17 00:00:00 2001 From: axisrow Date: Fri, 21 Aug 2026 08:30:41 +0700 Subject: [PATCH] feat: support resumable incremental manifests --- src/wordstat/cli.py | 13 ++++++- src/wordstat/collector.py | 81 +++++++++++++++++++++++++++++---------- src/wordstat/models.py | 9 +++++ src/wordstat/storage.py | 46 +++++++++++++++++++++- tests/test_storage.py | 64 ++++++++++++++++++++++++++++++- 5 files changed, 187 insertions(+), 26 deletions(-) diff --git a/src/wordstat/cli.py b/src/wordstat/cli.py index b607382..ee2aabe 100644 --- a/src/wordstat/cli.py +++ b/src/wordstat/cli.py @@ -74,6 +74,13 @@ def _read_phrases_file(path: Path) -> str: default=False, help="Keep each downloaded CSV as .csv instead of discarding it after conversion.", ) +@click.option( + "--resume", + "resume_directory", + type=click.Path(path_type=Path, file_okay=False, exists=True), + default=None, + help="Explicit existing run directory to complete; its phrase and region must match.", +) @click.pass_context def collect( ctx: click.Context, @@ -84,6 +91,7 @@ def collect( cdp_url: str, timeout_seconds: float, keep_raw: bool, + resume_directory: Path | None, ) -> None: """Collect all MVP Wordstat reports for one or more PHRASE as Parquet datasets. @@ -103,7 +111,10 @@ def collect( keep_raw=keep_raw, ) try: - batch = asyncio.run(collector.collect_many(phrases, region=region)) + if resume_directory is None: + batch = asyncio.run(collector.collect_many(phrases, region=region)) + else: + batch = asyncio.run(collector.collect_many(phrases, region=region, resume_directory=resume_directory)) # Only domain errors become friendly messages; an unexpected ValueError # from a dependency should keep its traceback instead of being reworded. except WordstatError as error: diff --git a/src/wordstat/collector.py b/src/wordstat/collector.py index 61fa37c..23ee413 100644 --- a/src/wordstat/collector.py +++ b/src/wordstat/collector.py @@ -23,11 +23,12 @@ BatchCollectionResult, CollectionManifest, CollectionResult, + CollectionStatus, ExportSummary, PhraseFailure, WordstatView, ) -from wordstat.storage import create_run_directory, finalize_raw, write_manifest +from wordstat.storage import create_run_directory, finalize_raw, validate_resume_directory, write_manifest WORDSTAT_URL = "https://wordstat.yandex.ru/" QUERY_SELECTOR = 'input[placeholder="Введите слово или словосочетание"]' @@ -73,10 +74,12 @@ def __init__( self.timeout_seconds = timeout_seconds self.keep_raw = keep_raw - async def collect(self, phrase: str, region: str = "Россия") -> CollectionResult: + async def collect( + self, phrase: str, region: str = "Россия", resume_directory: Path | None = None + ) -> CollectionResult: """Collect popular, related, dynamics and regional reports for one phrase.""" - batch = await self.collect_many([phrase], region=region) + batch = await self.collect_many([phrase], region=region, resume_directory=resume_directory) if batch.failures: raise batch.failures[0].error if not batch.results: @@ -87,7 +90,9 @@ async def collect(self, phrase: str, region: str = "Россия") -> Collection raise InvalidRequestError("Collecting the phrase produced neither a result nor a failure") return batch.results[0] - async def collect_many(self, phrases: list[str], region: str = "Россия") -> BatchCollectionResult: + async def collect_many( + self, phrases: list[str], region: str = "Россия", resume_directory: Path | None = None + ) -> BatchCollectionResult: """Collect reports for several phrases inside a single browser session. A failure on one phrase is recorded and does not stop the remaining @@ -106,6 +111,8 @@ async def collect_many(self, phrases: list[str], region: str = "Россия") - cleaned_phrases = [phrase.strip() for phrase in phrases] if not all(cleaned_phrases): raise InvalidRequestError("The search phrase must not be empty") + if resume_directory is not None and len(cleaned_phrases) != 1: + raise InvalidRequestError("--resume can only be used with one search phrase") results: list[CollectionResult] = [] failures: list[PhraseFailure] = [] @@ -164,14 +171,14 @@ def _mark_region_ready() -> None: # _collect_one for why the region control can't be # re-selected once a phrase's view loop has run (and # doesn't need to be after that). + collect_kwargs = { + "set_region": not region_ready, + "on_region_applied": _mark_region_ready, + } + if resume_directory is not None: + collect_kwargs["resume_directory"] = resume_directory result = await self._collect_one( - page, - session, - downloads_path, - phrase, - region, - set_region=not region_ready, - on_region_applied=_mark_region_ready, + page, session, downloads_path, phrase, region, **collect_kwargs ) results.append(result) except AuthenticationRequiredError as error: @@ -206,6 +213,7 @@ async def _collect_one( downloads_path: Path, phrase: str, region: str, + resume_directory: Path | None = None, set_region: bool = True, on_region_applied: Callable[[], None] | None = None, ) -> CollectionResult: @@ -216,7 +224,12 @@ async def _collect_one( # silently stopped working. await self._assert_authenticated(page) - run_directory = create_run_directory(self.output_root, phrase) + if resume_directory is None: + run_directory = create_run_directory(self.output_root, phrase) + existing_manifest = None + else: + run_directory = resume_directory + existing_manifest = validate_resume_directory(run_directory, phrase, region) # In a batch, the previous phrase's table can still be sitting in the # DOM when _set_phrase's own waits (new `words=` in the URL, download # button present) are satisfied — those don't check the table itself. @@ -252,8 +265,26 @@ async def _collect_one( required=False, ) - exports = [] + exports = list(existing_manifest.exports) if existing_manifest is not None else [] + manifest_path = run_directory / "manifest.json" + if existing_manifest is None: + manifest = CollectionManifest( + phrase=phrase, + region=region, + created_at=datetime.now(UTC), + source_url=await page.get_url(), + exports=[], + status=CollectionStatus.INCOMPLETE, + missing_views=list(WordstatView), + ) + write_manifest(manifest_path, manifest) + else: + manifest = existing_manifest + + completed_views = {export.view for export in exports} for view, selector in VIEW_SELECTORS.items(): + if view in completed_views: + continue await self._select_view(page, selector) source = await self._download_current_view(page, session, downloads_path) try: @@ -282,15 +313,23 @@ async def _collect_one( dtypes=dtypes, ) ) - - manifest = CollectionManifest( - phrase=phrase, - region=region, - created_at=datetime.now(UTC), - source_url=await page.get_url(), - exports=exports, + completed_views.add(view) + manifest = manifest.model_copy( + update={ + "exports": exports, + "status": CollectionStatus.INCOMPLETE, + "missing_views": [candidate for candidate in WordstatView if candidate not in completed_views], + } + ) + write_manifest(manifest_path, manifest) + + manifest = manifest.model_copy( + update={ + "exports": exports, + "status": CollectionStatus.COMPLETE, + "missing_views": [], + } ) - manifest_path = run_directory / "manifest.json" write_manifest(manifest_path, manifest) return CollectionResult( run_directory=run_directory, diff --git a/src/wordstat/models.py b/src/wordstat/models.py index 7394011..0b94dab 100644 --- a/src/wordstat/models.py +++ b/src/wordstat/models.py @@ -16,6 +16,13 @@ class WordstatView(StrEnum): REGIONS = "regions" +class CollectionStatus(StrEnum): + """Whether all four reports have been written to a run directory.""" + + COMPLETE = "complete" + INCOMPLETE = "incomplete" + + class CsvDataset(BaseModel): """A parsed CSV export with headers preserved exactly as Wordstat provided them. @@ -53,6 +60,8 @@ class CollectionManifest(BaseModel): created_at: datetime source_url: str exports: list[ExportSummary] + status: CollectionStatus = CollectionStatus.COMPLETE + missing_views: list[WordstatView] = Field(default_factory=list) class CollectionResult(BaseModel): diff --git a/src/wordstat/storage.py b/src/wordstat/storage.py index 86ebea3..cd0f988 100644 --- a/src/wordstat/storage.py +++ b/src/wordstat/storage.py @@ -1,9 +1,13 @@ """Run-directory and manifest handling.""" +import json +import os import re +import tempfile from datetime import UTC, datetime from pathlib import Path +from wordstat.errors import InvalidRequestError from wordstat.models import CollectionManifest, WordstatView @@ -45,8 +49,46 @@ def finalize_raw(source: Path, run_directory: Path, view: WordstatView, keep_raw def write_manifest(path: Path, manifest: CollectionManifest) -> None: - """Write reproducibility metadata in stable UTF-8 JSON.""" + """Atomically write reproducibility metadata in stable UTF-8 JSON.""" # pydantic emits UTF-8 without escaping non-ASCII, so the Cyrillic stays # readable without a round-trip through the json module. - path.write_text(manifest.model_dump_json(indent=2) + "\n", encoding="utf-8") + path.parent.mkdir(parents=True, exist_ok=True) + temporary: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="w", encoding="utf-8", dir=path.parent, prefix=f".{path.name}.", suffix=".tmp", delete=False + ) as handle: + temporary = Path(handle.name) + handle.write(manifest.model_dump_json(indent=2) + "\n") + handle.flush() + os.fsync(handle.fileno()) + os.replace(temporary, path) + temporary = None + finally: + if temporary is not None: + temporary.unlink(missing_ok=True) + + +def read_manifest(path: Path) -> CollectionManifest: + """Load a run manifest, preserving a useful domain error for bad input.""" + + try: + return CollectionManifest.model_validate(json.loads(path.read_text(encoding="utf-8"))) + except (OSError, ValueError) as error: + raise InvalidRequestError(f"Cannot read manifest {path}: {error}") from error + + +def validate_resume_directory(run_directory: Path, phrase: str, region: str) -> CollectionManifest: + """Ensure an explicitly selected run belongs to this exact request.""" + + manifest_path = run_directory / "manifest.json" + if not run_directory.is_dir() or not manifest_path.is_file(): + raise InvalidRequestError(f"Resume directory has no manifest: {run_directory}") + manifest = read_manifest(manifest_path) + if manifest.phrase != phrase or manifest.region != region: + raise InvalidRequestError( + f"Resume manifest request mismatch: expected phrase={phrase!r}, region={region!r}; " + f"found phrase={manifest.phrase!r}, region={manifest.region!r}" + ) + return manifest diff --git a/tests/test_storage.py b/tests/test_storage.py index 39ee787..6851407 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -1,8 +1,18 @@ from datetime import UTC, datetime from pathlib import Path -from wordstat.models import CollectionManifest, ExportSummary, WordstatView -from wordstat.storage import create_run_directory, finalize_raw, slugify, write_manifest +import pytest + +from wordstat.errors import InvalidRequestError +from wordstat.models import CollectionManifest, CollectionStatus, ExportSummary, WordstatView +from wordstat.storage import ( + create_run_directory, + finalize_raw, + read_manifest, + slugify, + validate_resume_directory, + write_manifest, +) def test_create_run_directory_is_unique_and_keeps_cyrillic(tmp_path: Path) -> None: @@ -42,6 +52,56 @@ def test_write_manifest_preserves_cyrillic_metadata(tmp_path: Path) -> None: assert '"phrase": "ремонт квартир"' in path.read_text(encoding="utf-8") +def test_write_manifest_does_not_corrupt_existing_file_if_replace_fails(tmp_path: Path, monkeypatch) -> None: + path = tmp_path / "manifest.json" + path.write_text("old manifest", encoding="utf-8") + manifest = CollectionManifest( + phrase="новый", region="Россия", created_at=datetime.now(UTC), source_url="url", exports=[] + ) + + def fail_replace(source, destination): + raise OSError("simulated interruption") + + monkeypatch.setattr("wordstat.storage.os.replace", fail_replace) + with pytest.raises(OSError, match="simulated interruption"): + write_manifest(path, manifest) + assert path.read_text(encoding="utf-8") == "old manifest" + + +def test_manifest_distinguishes_incomplete_run_from_complete(tmp_path: Path) -> None: + path = tmp_path / "manifest.json" + partial = CollectionManifest( + phrase="чай", region="Россия", created_at=datetime.now(UTC), source_url="url", exports=[], + status=CollectionStatus.INCOMPLETE, missing_views=list(WordstatView), + ) + write_manifest(path, partial) + loaded = read_manifest(path) + assert loaded.status is CollectionStatus.INCOMPLETE + assert loaded.missing_views == list(WordstatView) + + complete = partial.model_copy(update={"status": CollectionStatus.COMPLETE, "missing_views": []}) + write_manifest(path, complete) + loaded = read_manifest(path) + assert loaded.status is CollectionStatus.COMPLETE + assert loaded.missing_views == [] + + +def test_resume_validation_accepts_same_request_and_rejects_other_request(tmp_path: Path) -> None: + run = tmp_path / "run" + run.mkdir() + manifest = CollectionManifest( + phrase="чай", region="Москва", created_at=datetime.now(UTC), source_url="url", exports=[], + status=CollectionStatus.INCOMPLETE, missing_views=list(WordstatView), + ) + write_manifest(run / "manifest.json", manifest) + + assert validate_resume_directory(run, "чай", "Москва").phrase == "чай" + with pytest.raises(InvalidRequestError, match="mismatch"): + validate_resume_directory(run, "кофе", "Москва") + with pytest.raises(InvalidRequestError, match="mismatch"): + validate_resume_directory(run, "чай", "Россия") + + def test_finalize_raw_removes_the_download_by_default(tmp_path: Path) -> None: source = tmp_path / "wordstat-export.csv" source.write_text("Запрос;Показов\n", encoding="cp1251")