feat: incremental manifests and resumable runs

src/wordstat/cli.py

@@ -74,6 +74,13 @@ def _read_phrases_file(path: Path) -> str:
     default=False,
     help="Keep each downloaded CSV as <view>.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:

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,

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):

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

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")