From 4a80e7d56f8310127f991dc3db7bcb9cb41b01e6 Mon Sep 17 00:00:00 2001 From: axisrow Date: Fri, 21 Aug 2026 09:01:42 +0700 Subject: [PATCH] feat: support incremental Wordstat manifests --- src/wordstat/cli.py | 14 +++++- src/wordstat/collector.py | 61 +++++++++++++++++------- src/wordstat/models.py | 19 +++++++- src/wordstat/storage.py | 53 ++++++++++++++++++++- tests/test_collector_batch.py | 74 +++++++++++++++++++++++++++- tests/test_storage.py | 90 ++++++++++++++++++++++++++++------- 6 files changed, 272 insertions(+), 39 deletions(-) diff --git a/src/wordstat/cli.py b/src/wordstat/cli.py index b607382..be0752b 100644 --- a/src/wordstat/cli.py +++ b/src/wordstat/cli.py @@ -74,6 +74,12 @@ 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-run", + type=click.Path(path_type=Path, exists=True, file_okay=False), + default=None, + help="Explicit partial run directory to complete; its phrase and region must match.", +) @click.pass_context def collect( ctx: click.Context, @@ -84,6 +90,7 @@ def collect( cdp_url: str, timeout_seconds: float, keep_raw: bool, + resume_run: Path | None, ) -> None: """Collect all MVP Wordstat reports for one or more PHRASE as Parquet datasets. @@ -95,6 +102,8 @@ def collect( phrases = resolve_phrases(phrase, phrases_file) if not phrases: raise click.ClickException("At least one search phrase is required") + if resume_run is not None and len(phrases) != 1: + raise click.ClickException("--resume-run can only be used with one search phrase") collector = WordstatCollector( cdp_url=cdp_url, @@ -103,7 +112,10 @@ def collect( keep_raw=keep_raw, ) try: - batch = asyncio.run(collector.collect_many(phrases, region=region)) + collect_many_kwargs = {"region": region} + if resume_run is not None: + collect_many_kwargs["resume_run_directory"] = resume_run + batch = asyncio.run(collector.collect_many(phrases, **collect_many_kwargs)) # 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..59753e4 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, load_resume_manifest, 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_run_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_run_directory=resume_run_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_run_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_run_directory is not None and len(cleaned_phrases) != 1: + raise InvalidRequestError("--resume-run can only be used with one search phrase") results: list[CollectionResult] = [] failures: list[PhraseFailure] = [] @@ -164,15 +171,20 @@ 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). - result = await self._collect_one( + collect_one_args = ( page, session, downloads_path, phrase, region, - set_region=not region_ready, - on_region_applied=_mark_region_ready, ) + collect_one_kwargs = { + "set_region": not region_ready, + "on_region_applied": _mark_region_ready, + } + if resume_run_directory is not None: + collect_one_kwargs["resume_run_directory"] = resume_run_directory + result = await self._collect_one(*collect_one_args, **collect_one_kwargs) results.append(result) except AuthenticationRequiredError as error: # The session itself is gone: every remaining phrase @@ -208,6 +220,7 @@ async def _collect_one( region: str, set_region: bool = True, on_region_applied: Callable[[], None] | None = None, + resume_run_directory: Path | None = None, ) -> CollectionResult: # Checked once before the batch starts (collect_many), but a session # can lose authentication mid-batch (e.g. Yandex logs it out); check @@ -216,7 +229,13 @@ async def _collect_one( # silently stopped working. await self._assert_authenticated(page) - run_directory = create_run_directory(self.output_root, phrase) + if resume_run_directory is None: + run_directory = create_run_directory(self.output_root, phrase) + manifest_path = run_directory / "manifest.json" + manifest: CollectionManifest | None = None + else: + run_directory = resume_run_directory + manifest_path, manifest = load_resume_manifest(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 +271,20 @@ async def _collect_one( required=False, ) - exports = [] + if manifest is None: + manifest = CollectionManifest( + phrase=phrase, + region=region, + created_at=datetime.now(UTC), + source_url=await page.get_url(), + exports=[], + ) + write_manifest(manifest_path, manifest) + + exported_views = {export.view for export in manifest.exports} for view, selector in VIEW_SELECTORS.items(): + if view in exported_views: + continue await self._select_view(page, selector) source = await self._download_current_view(page, session, downloads_path) try: @@ -273,7 +304,7 @@ async def _collect_one( if source.exists(): source.replace(run_directory / source.name) raise - exports.append( + manifest.exports.append( ExportSummary( view=view, file=data_path.name, @@ -282,15 +313,9 @@ async def _collect_one( dtypes=dtypes, ) ) + write_manifest(manifest_path, manifest) - manifest = CollectionManifest( - phrase=phrase, - region=region, - created_at=datetime.now(UTC), - source_url=await page.get_url(), - exports=exports, - ) - manifest_path = run_directory / "manifest.json" + manifest.status = CollectionStatus.COMPLETE 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..3fcec65 100644 --- a/src/wordstat/models.py +++ b/src/wordstat/models.py @@ -4,7 +4,7 @@ from enum import StrEnum from pathlib import Path -from pydantic import BaseModel, ConfigDict, Field +from pydantic import BaseModel, ConfigDict, Field, model_validator class WordstatView(StrEnum): @@ -16,6 +16,13 @@ class WordstatView(StrEnum): REGIONS = "regions" +class CollectionStatus(StrEnum): + """Whether every required Wordstat view has been collected.""" + + IN_PROGRESS = "in_progress" + COMPLETE = "complete" + + class CsvDataset(BaseModel): """A parsed CSV export with headers preserved exactly as Wordstat provided them. @@ -53,6 +60,16 @@ class CollectionManifest(BaseModel): created_at: datetime source_url: str exports: list[ExportSummary] + status: CollectionStatus = CollectionStatus.IN_PROGRESS + + @model_validator(mode="after") + def exports_have_unique_views(self) -> "CollectionManifest": + """Reject a malformed resume manifest before it can duplicate output.""" + + views = [export.view for export in self.exports] + if len(views) != len(set(views)): + raise ValueError("Manifest contains duplicate view exports") + return self class CollectionResult(BaseModel): diff --git a/src/wordstat/storage.py b/src/wordstat/storage.py index 86ebea3..67ccc18 100644 --- a/src/wordstat/storage.py +++ b/src/wordstat/storage.py @@ -1,9 +1,14 @@ """Run-directory and manifest handling.""" +import os import re +import tempfile from datetime import UTC, datetime from pathlib import Path +from pydantic import ValidationError + +from wordstat.errors import InvalidRequestError from wordstat.models import CollectionManifest, WordstatView @@ -45,8 +50,52 @@ 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. + + Incremental collection rewrites the manifest after every successful view. + Writing a temporary file in the destination directory and replacing the + old file only after the write completes prevents an interrupted rewrite + from leaving a truncated manifest behind. + """ # 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") + descriptor, temporary_name = tempfile.mkstemp( + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + ) + temporary_path = Path(temporary_name) + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as temporary_file: + temporary_file.write(manifest.model_dump_json(indent=2) + "\n") + temporary_file.flush() + os.fsync(temporary_file.fileno()) + os.replace(temporary_path, path) + finally: + temporary_path.unlink(missing_ok=True) + + +def load_resume_manifest(run_directory: Path, phrase: str, region: str) -> tuple[Path, CollectionManifest]: + """Load an explicitly requested partial run after checking its identity. + + A run directory is never selected implicitly: this function only opens the + exact directory the caller supplied, and refuses to mix its Parquet files + with a different phrase or region. + """ + + manifest_path = run_directory / "manifest.json" + if not run_directory.is_dir(): + raise InvalidRequestError(f"Resume run directory does not exist: {run_directory}") + if not manifest_path.is_file(): + raise InvalidRequestError(f"Resume run directory has no manifest.json: {run_directory}") + try: + manifest = CollectionManifest.model_validate_json(manifest_path.read_text(encoding="utf-8")) + except (OSError, ValidationError) as error: + raise InvalidRequestError(f"Cannot read resume manifest: {manifest_path}") from error + if manifest.phrase != phrase or manifest.region != region: + raise InvalidRequestError( + "Resume run identity does not match the requested phrase and region " + f"({manifest.phrase!r}, {manifest.region!r})" + ) + return manifest_path, manifest diff --git a/tests/test_collector_batch.py b/tests/test_collector_batch.py index 88fa34f..dde4dd3 100644 --- a/tests/test_collector_batch.py +++ b/tests/test_collector_batch.py @@ -14,7 +14,8 @@ import wordstat.collector as collector_module from wordstat.collector import WordstatCollector from wordstat.errors import AuthenticationRequiredError, InvalidRequestError, PhraseEntryError -from wordstat.models import CollectionManifest, CollectionResult +from wordstat.models import CollectionManifest, CollectionResult, CollectionStatus, ExportSummary, WordstatView +from wordstat.storage import write_manifest def _fake_result(tmp_path, phrase) -> CollectionResult: @@ -455,3 +456,74 @@ async def run(): # Must raise the original RuntimeError, not a FileNotFoundError from the # rescue trying to move an already-moved file. asyncio.run(run()) + + run_directory = next(tmp_path.glob("runs/*")) + manifest = CollectionManifest.model_validate_json((run_directory / "manifest.json").read_text(encoding="utf-8")) + assert manifest.status is CollectionStatus.IN_PROGRESS + assert [export.view for export in manifest.exports] == [WordstatView.TOP_POPULAR] + + +def test_collect_one_resume_keeps_existing_exports_and_completes_missing_views(monkeypatch, tmp_path): + """A resume uses the supplied partial run and never rewrites its completed view.""" + + _patch_common(monkeypatch) + run_directory = tmp_path / "existing-run" + run_directory.mkdir() + existing_file = run_directory / "top_popular.parquet" + existing_file.write_text("original top popular dataset", encoding="utf-8") + write_manifest( + run_directory / "manifest.json", + CollectionManifest( + phrase="тест", + region="Россия", + created_at=datetime(2026, 8, 20, 12, 0, tzinfo=UTC), + source_url="https://wordstat.yandex.ru/?words=тест", + exports=[ + ExportSummary( + view=WordstatView.TOP_POPULAR, + file=existing_file.name, + raw_file=None, + row_count=1, + dtypes={"Запрос": "string"}, + ) + ], + ), + ) + downloads_path = tmp_path / "downloads" + downloads_path.mkdir() + downloaded_views = [] + + async def fake_select_view(self, page, selector): + downloaded_views.append(selector) + + async def fake_download(self, page, session, dl_path): + source = dl_path / f"export-{len(downloaded_views)}.csv" + _write_view_csv(source, "тест") + return source + + def fake_write_dataset(dataset, destination): + path = destination / f"{dataset.view.value}.parquet" + path.write_text(f"new {dataset.view.value} dataset", encoding="utf-8") + return path, {"Запрос": "string"} + + monkeypatch.setattr(WordstatCollector, "_select_view", fake_select_view) + monkeypatch.setattr(WordstatCollector, "_download_current_view", fake_download) + monkeypatch.setattr(collector_module, "write_dataset", fake_write_dataset) + + collector = WordstatCollector("cdp", tmp_path) + result = asyncio.run( + collector._collect_one( + _FakePage(), + _FakeSession(), + downloads_path, + "тест", + "Россия", + resume_run_directory=run_directory, + ) + ) + + assert len(downloaded_views) == 3 + assert existing_file.read_text(encoding="utf-8") == "original top popular dataset" + assert result.run_directory == run_directory + assert result.manifest.status is CollectionStatus.COMPLETE + assert [export.view for export in result.manifest.exports] == list(WordstatView) diff --git a/tests/test_storage.py b/tests/test_storage.py index 39ee787..e505a81 100644 --- a/tests/test_storage.py +++ b/tests/test_storage.py @@ -1,8 +1,25 @@ 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 + +import wordstat.storage as storage_module +from wordstat.errors import InvalidRequestError +from wordstat.models import CollectionManifest, CollectionStatus, ExportSummary, WordstatView +from wordstat.storage import create_run_directory, finalize_raw, load_resume_manifest, slugify, write_manifest + + +def _manifest( + phrase: str = "ремонт квартир", region: str = "Москва", status: CollectionStatus = CollectionStatus.IN_PROGRESS +) -> CollectionManifest: + return CollectionManifest( + phrase=phrase, + region=region, + created_at=datetime(2026, 8, 20, 12, 0, tzinfo=UTC), + source_url="https://wordstat.yandex.ru/?words=test", + exports=[], + status=status, + ) def test_create_run_directory_is_unique_and_keeps_cyrillic(tmp_path: Path) -> None: @@ -21,20 +38,15 @@ def test_slugify_uses_a_safe_fallback_for_symbols() -> None: def test_write_manifest_preserves_cyrillic_metadata(tmp_path: Path) -> None: path = tmp_path / "manifest.json" - manifest = CollectionManifest( - phrase="ремонт квартир", - region="Москва", - created_at=datetime(2026, 8, 20, 12, 0, tzinfo=UTC), - source_url="https://wordstat.yandex.ru/?words=test", - exports=[ - ExportSummary( - view=WordstatView.TOP_POPULAR, - file="top_popular.parquet", - raw_file=None, - row_count=1, - dtypes={"Запрос": "string"}, - ) - ], + manifest = _manifest() + manifest.exports.append( + ExportSummary( + view=WordstatView.TOP_POPULAR, + file="top_popular.parquet", + raw_file=None, + row_count=1, + dtypes={"Запрос": "string"}, + ) ) write_manifest(path, manifest) @@ -42,6 +54,52 @@ def test_write_manifest_preserves_cyrillic_metadata(tmp_path: Path) -> None: assert '"phrase": "ремонт квартир"' in path.read_text(encoding="utf-8") +def test_write_manifest_keeps_the_previous_file_if_atomic_replace_fails(monkeypatch, tmp_path: Path) -> None: + path = tmp_path / "manifest.json" + original = _manifest(phrase="старый запрос") + write_manifest(path, original) + previous_contents = path.read_text(encoding="utf-8") + + def fail_replace(source: Path, destination: Path) -> None: + raise OSError("simulated interruption before replacement") + + monkeypatch.setattr(storage_module.os, "replace", fail_replace) + + with pytest.raises(OSError, match="simulated interruption"): + write_manifest(path, _manifest(phrase="новый запрос")) + + assert path.read_text(encoding="utf-8") == previous_contents + assert list(tmp_path.glob(".manifest.json.*.tmp")) == [] + + +def test_partial_manifest_is_explicitly_distinct_from_a_complete_manifest(tmp_path: Path) -> None: + partial_path = tmp_path / "partial.json" + complete_path = tmp_path / "complete.json" + write_manifest(partial_path, _manifest()) + write_manifest(complete_path, _manifest(status=CollectionStatus.COMPLETE)) + + partial = CollectionManifest.model_validate_json(partial_path.read_text(encoding="utf-8")) + complete = CollectionManifest.model_validate_json(complete_path.read_text(encoding="utf-8")) + + assert partial.status is CollectionStatus.IN_PROGRESS + assert complete.status is CollectionStatus.COMPLETE + + +@pytest.mark.parametrize( + ("phrase", "region", "message"), + [("другая фраза", "Москва", "phrase and region"), ("ремонт квартир", "Россия", "phrase and region")], +) +def test_load_resume_manifest_rejects_a_different_query_identity( + tmp_path: Path, phrase: str, region: str, message: str +) -> None: + run_directory = tmp_path / "run" + run_directory.mkdir() + write_manifest(run_directory / "manifest.json", _manifest()) + + with pytest.raises(InvalidRequestError, match=message): + load_resume_manifest(run_directory, phrase, region) + + 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")