pivot 0.2.0__tar.gz → 0.2.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (107) hide show
  1. {pivot-0.2.0 → pivot-0.2.1}/PKG-INFO +1 -1
  2. {pivot-0.2.0 → pivot-0.2.1}/pyproject.toml +1 -1
  3. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/checkout.py +30 -6
  4. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/verify.py +27 -2
  5. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/executor/worker.py +7 -1
  6. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/import_artifact.py +38 -20
  7. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/remote/sync.py +38 -2
  8. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/cache.py +257 -2
  9. {pivot-0.2.0 → pivot-0.2.1}/LICENSE +0 -0
  10. {pivot-0.2.0 → pivot-0.2.1}/README.md +0 -0
  11. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/__init__.py +0 -0
  12. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/ast_utils.py +0 -0
  13. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/AGENTS.md +0 -0
  14. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/CLAUDE.md +0 -0
  15. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/__init__.py +0 -0
  16. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/_run_common.py +0 -0
  17. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/check_ignore.py +0 -0
  18. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/commit.py +0 -0
  19. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/completion.py +0 -0
  20. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/config.py +0 -0
  21. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/console.py +0 -0
  22. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/dag.py +0 -0
  23. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/data.py +0 -0
  24. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/decorators.py +0 -0
  25. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/doctor.py +0 -0
  26. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/errors.py +0 -0
  27. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/fingerprint.py +0 -0
  28. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/helpers.py +0 -0
  29. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/history.py +0 -0
  30. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/import_cmd.py +0 -0
  31. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/init.py +0 -0
  32. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/list.py +0 -0
  33. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/metrics.py +0 -0
  34. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/params.py +0 -0
  35. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/plots.py +0 -0
  36. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/remote.py +0 -0
  37. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/repro.py +0 -0
  38. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/run.py +0 -0
  39. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/status.py +0 -0
  40. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/targets.py +0 -0
  41. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/track.py +0 -0
  42. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/cli/update.py +0 -0
  43. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/compose.py +0 -0
  44. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/config/__init__.py +0 -0
  45. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/config/io.py +0 -0
  46. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/config/models.py +0 -0
  47. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/dag/__init__.py +0 -0
  48. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/dag/render.py +0 -0
  49. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/decorators.py +0 -0
  50. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/discovery.py +0 -0
  51. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/__init__.py +0 -0
  52. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/agent_rpc.py +0 -0
  53. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/engine.py +0 -0
  54. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/graph.py +0 -0
  55. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/scheduler.py +0 -0
  56. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/sinks.py +0 -0
  57. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/sources.py +0 -0
  58. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/types.py +0 -0
  59. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/watch.py +0 -0
  60. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/engine/worker_pool.py +0 -0
  61. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/exceptions.py +0 -0
  62. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/executor/AGENTS.md +0 -0
  63. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/executor/CLAUDE.md +0 -0
  64. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/executor/__init__.py +0 -0
  65. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/executor/commit.py +0 -0
  66. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/executor/core.py +0 -0
  67. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/explain.py +0 -0
  68. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/fingerprint.py +0 -0
  69. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/git.py +0 -0
  70. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/ignore.py +0 -0
  71. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/loaders.py +0 -0
  72. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/metrics.py +0 -0
  73. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/names.py +0 -0
  74. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/outputs.py +0 -0
  75. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/parameters.py +0 -0
  76. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/path_policy.py +0 -0
  77. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/project.py +0 -0
  78. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/py.typed +0 -0
  79. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/registry.py +0 -0
  80. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/remote/AGENTS.md +0 -0
  81. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/remote/__init__.py +0 -0
  82. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/remote/config.py +0 -0
  83. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/remote/git_archive.py +0 -0
  84. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/remote/github.py +0 -0
  85. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/remote/storage.py +0 -0
  86. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/run_history.py +0 -0
  87. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/show/__init__.py +0 -0
  88. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/show/common.py +0 -0
  89. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/show/data.py +0 -0
  90. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/show/metrics.py +0 -0
  91. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/show/params.py +0 -0
  92. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/show/plots.py +0 -0
  93. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/skip.py +0 -0
  94. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/stage_def.py +0 -0
  95. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/status.py +0 -0
  96. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/AGENTS.md +0 -0
  97. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/CLAUDE.md +0 -0
  98. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/__init__.py +0 -0
  99. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/artifact_lock.py +0 -0
  100. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/lock.py +0 -0
  101. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/presentation.py +0 -0
  102. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/restore.py +0 -0
  103. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/state.py +0 -0
  104. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/store.py +0 -0
  105. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/storage/track.py +0 -0
  106. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/types.py +0 -0
  107. {pivot-0.2.0 → pivot-0.2.1}/src/pivot/yaml_config.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: pivot
3
- Version: 0.2.0
3
+ Version: 0.2.1
4
4
  Summary: High-performance Python pipeline tool with automatic code change detection
5
5
  Author: Sami Jawhar
6
6
  Author-email: Sami Jawhar <sami@metr.org>
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "pivot"
3
- version = "0.2.0"
3
+ version = "0.2.1"
4
4
  description = "High-performance Python pipeline tool with automatic code change detection"
5
5
  requires-python = ">=3.13,<3.14"
6
6
  readme = "README.md"
@@ -1,6 +1,7 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import asyncio
4
+ import contextlib
4
5
  import enum
5
6
  import pathlib
6
7
  from typing import TYPE_CHECKING, Literal
@@ -12,7 +13,7 @@ from pivot.cli import completion
12
13
  from pivot.cli import decorators as cli_decorators
13
14
  from pivot.cli import helpers as cli_helpers
14
15
  from pivot.cli import targets as cli_targets
15
- from pivot.storage import cache, lock, track
16
+ from pivot.storage import cache, lock, state, track
16
17
  from pivot.types import HashInfo, is_dir_hash
17
18
 
18
19
  if TYPE_CHECKING:
@@ -96,6 +97,7 @@ def _restore_path_sync(
96
97
  "already exists" without --force). Cache misses return ("missing", name)
97
98
  instead of raising.
98
99
  """
100
+ keep_ignored = False
99
101
  if path.exists():
100
102
  match behavior:
101
103
  case CheckoutBehavior.ERROR:
@@ -106,12 +108,33 @@ def _restore_path_sync(
106
108
  + "Use --force to overwrite or --only-missing to skip existing files."
107
109
  )
108
110
  case CheckoutBehavior.SKIP_EXISTING:
109
- # For directories with manifests, don't skip - files inside may be missing.
110
- # Let restore_from_cache() handle it (does full directory restoration).
111
- # DirHash has "manifest" key, FileHash does not.
112
- is_directory = is_dir_hash(output_hash)
113
- if not is_directory:
111
+ # Existing file: skip. Existing real directory: fill in its missing files, or
112
+ # swap in the target when every differing file is pivot's own cached output,
113
+ # keeping the entries the directory hash skips. Never overwrite content the
114
+ # cache lacks (a local change).
115
+ if not is_dir_hash(output_hash) or path.is_symlink() or not path.is_dir():
114
116
  return ("skipped", path.name)
117
+ # Reuse hashes `pivot status` and `pivot track` recorded; never create the DB.
118
+ with (
119
+ state.StateDB(state_dir, readonly=True)
120
+ if state_dir is not None and (state_dir / "state.lmdb").exists()
121
+ else contextlib.nullcontext(None)
122
+ ) as state_db:
123
+ outcome = cache.restore_missing_in_directory(
124
+ path, output_hash, cache_dir, checkout_modes, state_db
125
+ )
126
+ if outcome["local_changes"]:
127
+ raise click.ClickException(
128
+ f"'{path.name}' has local changes that are not in the cache: "
129
+ + f"{', '.join(outcome['local_changes'])}. Use --force to overwrite."
130
+ )
131
+ # An interrupted restore may have staged kept entries beside the directory;
132
+ # the restore below puts them back, so take it even when nothing else differs.
133
+ if not outcome["replace"] and cache.staged_kept_entries(path) is None:
134
+ if outcome["unavailable"]:
135
+ return ("missing", path.name)
136
+ return ("restored" if outcome["restored"] else "skipped", path.name)
137
+ keep_ignored = True
115
138
  case CheckoutBehavior.FORCE:
116
139
  cache.remove_output(path)
117
140
  case _: # pyright: ignore[reportUnnecessaryComparison] - defensive for future enum values
@@ -123,6 +146,7 @@ def _restore_path_sync(
123
146
  cache_dir,
124
147
  checkout_modes=checkout_modes,
125
148
  state_dir=state_dir,
149
+ keep_ignored=keep_ignored,
126
150
  )
127
151
  if not success:
128
152
  return ("missing", path.name)
@@ -190,6 +190,29 @@ def _create_remote_if_needed(allow_missing: bool) -> remote_mod.S3Remote | None:
190
190
  return remote
191
191
 
192
192
 
193
+ def _status_from_explanations(
194
+ explanations: list[types.StageExplanation],
195
+ ) -> list[PipelineStatusInfo]:
196
+ """Convert rich explanations to verification status, retaining changed dep identities."""
197
+ statuses = list[PipelineStatusInfo]()
198
+ for explanation in explanations:
199
+ reason = explanation["reason"]
200
+ if reason == "Input dependencies changed" and explanation["dep_changes"]:
201
+ changed = ", ".join(
202
+ types.identity_key(change["identity"]) for change in explanation["dep_changes"]
203
+ )
204
+ reason = f"{reason}: {changed}"
205
+ statuses.append(
206
+ PipelineStatusInfo(
207
+ name=explanation["stage_name"],
208
+ status=(PipelineStatus.STALE if explanation["will_run"] else PipelineStatus.CACHED),
209
+ reason=reason,
210
+ upstream_stale=explanation["upstream_stale"],
211
+ )
212
+ )
213
+ return statuses
214
+
215
+
193
216
  def _verify_stages(
194
217
  pipeline_status: list[PipelineStatusInfo],
195
218
  cache_dir: pathlib.Path,
@@ -345,14 +368,16 @@ def verify(
345
368
 
346
369
  cache_dir = config.get_cache_dir()
347
370
 
348
- # Get pipeline status (uses default state directory internally)
349
- pipeline_status, _ = status_mod.get_pipeline_status(
371
+ # Get rich explanations so lock-only dependency mismatches identify the
372
+ # producer artifact in verification output.
373
+ explanations = status_mod.get_pipeline_explanations(
350
374
  stages_list,
351
375
  single_stage=False,
352
376
  all_stages=all_stages,
353
377
  pipeline=cli_helpers.get_pipeline(),
354
378
  allow_missing=allow_missing,
355
379
  )
380
+ pipeline_status = _status_from_explanations(explanations)
356
381
 
357
382
  if not pipeline_status:
358
383
  raise click.ClickException("No stages to verify.")
@@ -6,6 +6,7 @@ Must be module-level and picklable.
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
+ import codecs
9
10
  import collections
10
11
  import collections.abc as collections_abc
11
12
  import contextlib
@@ -1095,14 +1096,19 @@ class _QueueWriter:
1095
1096
 
1096
1097
  def _pipe_reader(self) -> None:
1097
1098
  assert self._read_fd is not None # Guaranteed by _ensure_pipe() before thread starts
1099
+ # One incremental decoder per pipe: a multi-byte character split across two reads
1100
+ # is held until its remaining bytes arrive instead of becoming U+FFFD.
1101
+ decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
1098
1102
  try:
1099
1103
  while True:
1100
1104
  data = os.read(self._read_fd, 8192)
1101
1105
  if not data:
1102
1106
  break
1103
- self.write(data.decode("utf-8", errors="replace"))
1107
+ self.write(decoder.decode(data))
1104
1108
  except OSError:
1105
1109
  pass
1110
+ # Flush a truncated trailing character so it is reported rather than dropped.
1111
+ self.write(decoder.decode(b"", final=True))
1106
1112
 
1107
1113
  def _ensure_pipe(self) -> None:
1108
1114
  if self._read_fd is not None:
@@ -1,8 +1,12 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import asyncio
4
+ import contextlib
4
5
  import functools
5
6
  import logging
7
+ import os
8
+ import pathlib
9
+ import tempfile
6
10
  from typing import TYPE_CHECKING, TypedDict, cast
7
11
 
8
12
  import aiohttp
@@ -16,7 +20,6 @@ from pivot.storage import cache, track
16
20
  from pivot.storage import lock as storage_lock
17
21
 
18
22
  if TYPE_CHECKING:
19
- import pathlib
20
23
  from collections.abc import Callable
21
24
 
22
25
  from pivot.types import DirManifestEntry, OutEntry, StorageLockData
@@ -39,6 +42,31 @@ class ImportResult(TypedDict):
39
42
  downloaded: bool
40
43
 
41
44
 
45
+ async def _download_verified(remote_url: str, file_hash: str, data_path: pathlib.Path) -> int:
46
+ """Download ``file_hash`` to ``data_path`` only if its content matches; return its size.
47
+
48
+ The file is downloaded to a temp file in the destination directory and moved over
49
+ ``data_path`` after its hash is checked, so a corrupt remote copy leaves any existing
50
+ file (the last good import) untouched.
51
+ """
52
+ fd, tmp_name = tempfile.mkstemp(
53
+ dir=data_path.parent, prefix=f".{data_path.name}.", suffix=".tmp"
54
+ )
55
+ os.close(fd)
56
+ tmp_path = pathlib.Path(tmp_name)
57
+ try:
58
+ await remote_storage.S3Remote(remote_url).download_file(file_hash, tmp_path)
59
+ actual_hash, _ = cache.hash_file(tmp_path)
60
+ if actual_hash != file_hash:
61
+ raise exceptions.RemoteError(f"Hash mismatch: expected {file_hash}, got {actual_hash}")
62
+ size = tmp_path.stat().st_size
63
+ os.replace(tmp_path, data_path)
64
+ finally:
65
+ with contextlib.suppress(FileNotFoundError):
66
+ tmp_path.unlink()
67
+ return size
68
+
69
+
42
70
  def _maybe_token(repo_url: str, token: str | None) -> str | None:
43
71
  if token is not None:
44
72
  return token
@@ -111,7 +139,13 @@ async def read_remote_config(
111
139
  remotes = {str(k): str(v) for k, v in remotes_data.items()}
112
140
 
113
141
  default_remote = config_data.get("default_remote")
114
- if not isinstance(default_remote, str):
142
+ if default_remote is not None and not isinstance(default_remote, str):
143
+ # Warn-first: the source config is a saved input outside this repo.
144
+ logger.warning(
145
+ "Remote config default_remote in %s must be a remote name, got %r; ignoring it",
146
+ repo_url,
147
+ default_remote,
148
+ )
115
149
  default_remote = None
116
150
  if default_remote is None:
117
151
  if len(remotes) == 1:
@@ -370,16 +404,8 @@ async def import_artifact(
370
404
  downloaded = False
371
405
  file_size = resolved["size"]
372
406
  if not no_download:
373
- remote = remote_storage.S3Remote(resolved["remote_url"])
374
- await remote.download_file(resolved["hash"], data_path)
407
+ file_size = await _download_verified(resolved["remote_url"], resolved["hash"], data_path)
375
408
  downloaded = True
376
- actual_hash, _ = cache.hash_file(data_path)
377
- if actual_hash != resolved["hash"]:
378
- data_path.unlink(missing_ok=True)
379
- raise exceptions.RemoteError(
380
- f"Downloaded file hash mismatch: expected {resolved['hash']}, got {actual_hash}"
381
- )
382
- file_size = data_path.stat().st_size
383
409
 
384
410
  source = track.ImportSource(
385
411
  repo=repo_url,
@@ -456,16 +482,8 @@ async def update_import(pvt_path: pathlib.Path, *, new_rev: str | None = None) -
456
482
 
457
483
  file_size = resolved["size"]
458
484
  if need_download:
459
- remote = remote_storage.S3Remote(resolved["remote_url"])
460
- await remote.download_file(resolved["hash"], data_path)
461
- actual_hash, _ = cache.hash_file(data_path)
462
- if actual_hash != resolved["hash"]:
463
- data_path.unlink(missing_ok=True)
464
- raise exceptions.RemoteError(
465
- f"Hash mismatch: expected {resolved['hash']}, got {actual_hash}"
466
- )
485
+ file_size = await _download_verified(resolved["remote_url"], resolved["hash"], data_path)
467
486
  downloaded = True
468
- file_size = data_path.stat().st_size
469
487
  elif data_path.exists():
470
488
  file_size = data_path.stat().st_size
471
489
 
@@ -180,6 +180,30 @@ def _get_file_hash_from_pvt(rel_path: str, proj_root: pathlib.Path) -> HashInfo
180
180
  return FileHash(hash=track_data["hash"])
181
181
 
182
182
 
183
+ def _get_tracked_hashes(proj_root: pathlib.Path) -> set[str]:
184
+ """Blob hashes of every `pivot track`-ed file or directory in the project.
185
+
186
+ Imports are left out: `pivot import` writes the file straight to the workspace, so its
187
+ blob never enters the local cache and push never sends it to this project's remote.
188
+ """
189
+ hashes = set[str]()
190
+ for pvt_data in track.discover_pvt_files(proj_root).values():
191
+ if track.is_import(pvt_data):
192
+ continue
193
+ hashes |= _extract_file_hashes_from_hash_info(track.pvt_to_hash_info(pvt_data))
194
+ return hashes
195
+
196
+
197
+ def _is_cached_stage_output(
198
+ identity: types.ArtifactIdentity, producer_info: RegistryStageInfo
199
+ ) -> bool:
200
+ """Whether a stage output is stored in the cache (Metrics are git-tracked, not cached)."""
201
+ return not any(
202
+ out.identity == identity and out.tag is types.ArtifactTag.METRIC
203
+ for out in producer_info["outs"]
204
+ )
205
+
206
+
183
207
  def get_target_hashes(
184
208
  targets: list[str],
185
209
  state_dir: pathlib.Path,
@@ -192,6 +216,7 @@ def get_target_hashes(
192
216
  hashes = set[str]()
193
217
  output_identities: dict[str, types.ArtifactIdentity] | None = None
194
218
  unresolved = list[str]()
219
+ tracked_hashes: set[str] | None = None
195
220
 
196
221
  for target in targets:
197
222
  original_target = target
@@ -224,8 +249,19 @@ def get_target_hashes(
224
249
  if types.identity_key(out_path) not in non_cached_paths:
225
250
  hashes |= _extract_file_hashes_from_hash_info(out_hash)
226
251
  if include_deps:
227
- for dep_hash in lock_data["dep_hashes"].values():
228
- hashes |= _extract_file_hashes_from_hash_info(dep_hash)
252
+ # Only deps whose blobs pivot caches (and so pushes) can be fetched:
253
+ # upstream stage outputs and `pivot track` inputs. Raw inputs
254
+ # committed to git and `pivot import` inputs are never in the remote.
255
+ for dep_identity, dep_hash in lock_data["dep_hashes"].items():
256
+ dep_blobs = _extract_file_hashes_from_hash_info(dep_hash)
257
+ producer_info = all_stages.get(dep_identity.producer)
258
+ if producer_info is not None:
259
+ if _is_cached_stage_output(dep_identity, producer_info):
260
+ hashes |= dep_blobs
261
+ continue
262
+ if tracked_hashes is None:
263
+ tracked_hashes = _get_tracked_hashes(proj_root)
264
+ hashes |= dep_blobs & tracked_hashes
229
265
  continue
230
266
  else:
231
267
  resolved_identity = types.ArtifactIdentity(stage_name, target_identity.key)
@@ -14,7 +14,7 @@ import shutil
14
14
  import stat
15
15
  import tempfile
16
16
  import time
17
- from typing import TYPE_CHECKING
17
+ from typing import TYPE_CHECKING, TypedDict
18
18
 
19
19
  import xxhash
20
20
 
@@ -36,9 +36,54 @@ XXHASH64_HEX_LENGTH = 16 # xxhash64 produces 64-bit hash = 16 hex characters
36
36
  _RESTORE_TEMP_PREFIX = ".pivot_restore_"
37
37
  _BACKUP_TEMP_PREFIX = ".pivot_backup_"
38
38
  _RESTORE_LOCK_PREFIX = ".pivot_restore_lock_"
39
+ _KEEP_STAGING_PREFIX = ".pivot_keep_"
39
40
  _RESTORE_TEMP_MAX_AGE = 3600 # 1 hour
40
41
 
41
42
 
43
+ def _keep_staging_path(path: pathlib.Path) -> pathlib.Path:
44
+ """Sibling that holds a directory's kept entries (see keep_ignored) while it is replaced.
45
+
46
+ It is named after the target and no cleanup removes it, so entries an interrupted
47
+ restore leaves there wait for the next restore of the same target to put them back.
48
+ """
49
+ logical_path = path.parent.resolve() / path.name
50
+ token = hashlib.sha256(str(logical_path).encode()).hexdigest()[:32]
51
+ return path.parent / f"{_KEEP_STAGING_PREFIX}{token}"
52
+
53
+
54
+ def _move_kept_entries(source: pathlib.Path, destination: pathlib.Path) -> None:
55
+ """Move every entry under source into destination, merging directories, then drop source.
56
+
57
+ Raises CacheRestoreError, leaving that entry and the rest in source, when a file or
58
+ link already exists at its destination.
59
+ """
60
+ with os.scandir(source) as entries:
61
+ for entry in entries:
62
+ target = destination / entry.name
63
+ if not os.path.lexists(target):
64
+ destination.mkdir(parents=True, exist_ok=True)
65
+ pathlib.Path(entry.path).rename(target)
66
+ elif (
67
+ entry.is_dir(follow_symlinks=False) and target.is_dir() and not target.is_symlink()
68
+ ):
69
+ _move_kept_entries(pathlib.Path(entry.path), target)
70
+ else:
71
+ raise exceptions.CacheRestoreError(
72
+ f"Cannot put back {entry.path}: {target} already exists. "
73
+ + "Move the entry by hand, then retry."
74
+ )
75
+ source.rmdir()
76
+
77
+
78
+ def staged_kept_entries(path: pathlib.Path) -> pathlib.Path | None:
79
+ """The directory where an interrupted restore of path left kept entries, if any.
80
+
81
+ Any later restore of path (restore_from_cache) puts them back.
82
+ """
83
+ keep_dir = _keep_staging_path(path)
84
+ return keep_dir if keep_dir.exists() else None
85
+
86
+
42
87
  def _get_lock_path(target_path: pathlib.Path, state_dir: pathlib.Path) -> pathlib.Path:
43
88
  """Get lock file path under state_dir/locks/ for a restore target.
44
89
 
@@ -605,6 +650,8 @@ def restore_from_cache(
605
650
  checkout_mode: CheckoutMode | None = None,
606
651
  checkout_modes: list[CheckoutMode] | None = None,
607
652
  state_dir: pathlib.Path | None = None,
653
+ *,
654
+ keep_ignored: bool = False,
608
655
  ) -> bool:
609
656
  """Restore file or directory from cache. Returns True if successful.
610
657
 
@@ -616,12 +663,16 @@ def restore_from_cache(
616
663
  checkout_modes: Ordered list of link modes to try with fallback on failure.
617
664
  state_dir: State directory for lock files (e.g. per-stage .pivot/ dir).
618
665
  Required for directory restores to prevent concurrent restore races.
666
+ keep_ignored: When replacing an existing directory, move the entries the directory
667
+ hash skips (build/, *.swp, ...) into the restored tree instead of deleting them.
668
+ They wait in a sibling staging directory during the swap. Entries an interrupted
669
+ restore left there are put into the restored tree by any later restore of path.
619
670
  """
620
671
  effective_modes = _resolve_checkout_modes(checkout_mode, checkout_modes)
621
672
 
622
673
  if is_dir_hash(output_hash):
623
674
  return _restore_directory_from_cache(
624
- path, output_hash, cache_dir, effective_modes, state_dir
675
+ path, output_hash, cache_dir, effective_modes, state_dir, keep_ignored=keep_ignored
625
676
  )
626
677
  return _restore_file_from_cache(path, output_hash, cache_dir, effective_modes)
627
678
 
@@ -648,6 +699,8 @@ def _restore_directory_from_cache(
648
699
  cache_dir: pathlib.Path,
649
700
  checkout_modes: list[CheckoutMode],
650
701
  state_dir: pathlib.Path | None = None,
702
+ *,
703
+ keep_ignored: bool = False,
651
704
  ) -> bool:
652
705
  """Restore directory from cache atomically with file locking."""
653
706
  cache_dir_path = get_cache_path(cache_dir, output_hash["hash"])
@@ -678,6 +731,23 @@ def _restore_directory_from_cache(
678
731
  lock_fd = os.open(lock_path, os.O_CREAT | os.O_RDWR, 0o644)
679
732
  fcntl.flock(lock_fd, fcntl.LOCK_EX)
680
733
 
734
+ # Entries of the existing directory that are not part of the output, to carry over.
735
+ # A restore that carries them first puts back what an interrupted one left staged in
736
+ # keep_dir, so they are carried like the rest. Any other restore leaves keep_dir alone
737
+ # and moves its contents into the new tree after the swap.
738
+ keep_dir = _keep_staging_path(path)
739
+ keep = list[str]()
740
+ if keep_ignored and path.is_dir() and not path.is_symlink():
741
+ if keep_dir.exists():
742
+ _move_kept_entries(keep_dir, path)
743
+ _, keep = _scan_entries(path, {entry["relpath"] for entry in output_hash["manifest"]})
744
+ if symlink_fast_path and (keep or keep_dir.exists()):
745
+ # A directory that keeps local entries cannot be a link to the cache tree.
746
+ symlink_fast_path = False
747
+ for entry in output_hash["manifest"]:
748
+ if not get_cache_path(cache_dir, entry["hash"]).exists():
749
+ return False
750
+
681
751
  # SYMLINK fast path: create symlink at temp location, then atomic rename
682
752
  if symlink_fast_path:
683
753
  temp_symlink = path.parent / f"{_RESTORE_TEMP_PREFIX}symlink_{secrets.token_hex(4)}"
@@ -748,6 +818,17 @@ def _restore_directory_from_cache(
748
818
  except OSError as e:
749
819
  logger.debug(f"chmod failed (may be NFS with root_squash): {e}")
750
820
 
821
+ # Entries the directory hash skips are not part of the output. Stage them in
822
+ # keep_dir, which no cleanup removes, and put them into the new tree after the
823
+ # swap. A crash at any point leaves each one in path, keep_dir or the new tree.
824
+ for relpath in keep:
825
+ source = path / relpath
826
+ if not os.path.lexists(source):
827
+ continue
828
+ staged = keep_dir / relpath
829
+ staged.parent.mkdir(parents=True, exist_ok=True)
830
+ source.rename(staged)
831
+
751
832
  # Atomic swap with rollback capability:
752
833
  # - Backup original first so we can restore on failure
753
834
  # - Replace atomically so path always has valid content
@@ -761,6 +842,9 @@ def _restore_directory_from_cache(
761
842
  temp_path.replace(path)
762
843
  swap_completed = True
763
844
 
845
+ if keep_dir.exists():
846
+ _move_kept_entries(keep_dir, path)
847
+
764
848
  try:
765
849
  _clear_path(backup_path)
766
850
  except OSError as e:
@@ -776,6 +860,15 @@ def _restore_directory_from_cache(
776
860
  except OSError as e:
777
861
  logger.debug(f"Failed to restore backup: {e}")
778
862
  backup_restore_failed = True
863
+ # Return staged entries to the original directory if the swap did not happen
864
+ if not swap_completed and keep_dir.exists() and path.is_dir() and not path.is_symlink():
865
+ try:
866
+ _move_kept_entries(keep_dir, path)
867
+ except (OSError, exceptions.CacheRestoreError) as e:
868
+ logger.warning(
869
+ f"Entries kept from {path} remain in {keep_dir} ({e}); "
870
+ + "the next restore of it puts them back"
871
+ )
779
872
  raise
780
873
  finally:
781
874
  # Clean up temp directory on any exit
@@ -807,6 +900,168 @@ def remove_output(path: pathlib.Path) -> None:
807
900
  _clear_path(path)
808
901
 
809
902
 
903
+ class MissingRestore(TypedDict):
904
+ """Outcome of restore_missing_in_directory."""
905
+
906
+ restored: int # missing files restored in place
907
+ unavailable: int # missing files whose blob is not in the cache
908
+ local_changes: list[str] # changed or extra files whose content the cache lacks
909
+ replace: bool # directory holds only cached content but not the target: replace it whole
910
+
911
+
912
+ def _is_cached(cache_dir: pathlib.Path, content_hash: str | None) -> bool:
913
+ return (
914
+ content_hash is not None
915
+ and is_valid_hash(content_hash)
916
+ and get_cache_path(cache_dir, content_hash).exists()
917
+ )
918
+
919
+
920
+ def _scan_entries(path: pathlib.Path, expected_paths: set[str]) -> tuple[list[str], list[str]]:
921
+ """Return (files and symlinks, entries the directory hash skips) under path, as relpaths.
922
+
923
+ Symlinked directories are not followed. A skipped name that is, or holds, a manifest
924
+ path is part of the output and is scanned like any other entry.
925
+ """
926
+ expected_dirs = {
927
+ str(parent)
928
+ for relpath in expected_paths
929
+ for parent in pathlib.PurePosixPath(relpath).parents
930
+ if str(parent) != "."
931
+ }
932
+ files = list[str]()
933
+ ignored = list[str]()
934
+
935
+ def _walk(directory: pathlib.Path) -> None:
936
+ with os.scandir(directory) as entries:
937
+ for entry in entries:
938
+ entry_path = pathlib.Path(entry.path)
939
+ relpath = str(entry_path.relative_to(path))
940
+ if (
941
+ _should_skip_entry(entry)
942
+ and relpath not in expected_paths
943
+ and relpath not in expected_dirs
944
+ ):
945
+ ignored.append(relpath)
946
+ elif entry.is_symlink() or entry.is_file():
947
+ files.append(relpath)
948
+ elif entry.is_dir():
949
+ _walk(entry_path)
950
+
951
+ _walk(path)
952
+ return files, ignored
953
+
954
+
955
+ def _is_same_file(file_stat: os.stat_result, cache_path: pathlib.Path) -> bool:
956
+ """Whether file_stat is the cache blob itself (a hardlink checkout). Cache blobs are
957
+ read-only and never rewritten, so sharing the inode proves the content without a read."""
958
+ try:
959
+ blob_stat = cache_path.stat()
960
+ except FileNotFoundError:
961
+ return False
962
+ return (file_stat.st_dev, file_stat.st_ino) == (blob_stat.st_dev, blob_stat.st_ino)
963
+
964
+
965
+ def restore_missing_in_directory(
966
+ path: pathlib.Path,
967
+ output_hash: DirHash,
968
+ cache_dir: pathlib.Path,
969
+ checkout_modes: list[CheckoutMode],
970
+ state_db: state_mod.StateDB | None = None,
971
+ ) -> MissingRestore:
972
+ """Restore a directory's missing target files without overwriting local changes.
973
+
974
+ When every existing file matches the target, the missing files are restored in place
975
+ and nothing else is touched. When some file differs from the target or is not in it,
976
+ only a whole-directory swap brings the target. It is safe only if every file it
977
+ removes is in the cache (pivot's own earlier output): then nothing is written and
978
+ ``replace`` asks the caller to swap with ``restore_from_cache(..., keep_ignored=True)``,
979
+ which carries over the entries the directory hash skips (build/, *.swp, ...). Otherwise
980
+ nothing is written, and the files the cache lacks (local edits, extra files) are returned.
981
+
982
+ A retained file is read only when nothing else proves its content: a symlink to the
983
+ target's cache path, a hardlink to the target's cache blob, or a matching ``state_db``
984
+ entry each settle it without a read.
985
+ """
986
+ lexical_base = pathlib.Path(os.path.abspath(path))
987
+ resolved_base = path.resolve()
988
+ expected_paths = set[str]()
989
+ missing = list[tuple[DirManifestEntry, pathlib.Path, pathlib.Path]]()
990
+ # relpath -> hash of the content on disk (None when it cannot be in the cache)
991
+ changed = dict[str, str | None]()
992
+
993
+ for entry in output_hash["manifest"]:
994
+ relpath = entry["relpath"]
995
+ expected_paths.add(relpath)
996
+ file_path = path / relpath
997
+ lexical_path = pathlib.Path(os.path.abspath(file_path))
998
+ if not lexical_path.is_relative_to(lexical_base):
999
+ raise exceptions.SecurityValidationError(
1000
+ f"Manifest contains path traversal: {relpath!r}"
1001
+ )
1002
+
1003
+ file_cache_path = get_cache_path(cache_dir, entry["hash"])
1004
+ if file_path.is_symlink():
1005
+ if not file_path.parent.resolve().is_relative_to(resolved_base):
1006
+ raise exceptions.SecurityValidationError(
1007
+ f"Manifest contains path traversal: {relpath!r}"
1008
+ )
1009
+ if file_path.resolve() == file_cache_path.resolve():
1010
+ if not file_cache_path.exists():
1011
+ missing.append((entry, file_path, file_cache_path))
1012
+ else:
1013
+ changed[relpath] = _get_symlink_cache_hash(file_path, cache_dir)
1014
+ continue
1015
+ if file_path.exists():
1016
+ if not file_path.is_file():
1017
+ changed[relpath] = None
1018
+ continue
1019
+ if _is_same_file(file_path.stat(), file_cache_path):
1020
+ continue
1021
+ file_hash, _ = hash_file(file_path, state_db)
1022
+ if file_hash != entry["hash"]:
1023
+ changed[relpath] = file_hash
1024
+ continue
1025
+ if not file_path.parent.resolve().is_relative_to(resolved_base):
1026
+ raise exceptions.SecurityValidationError(
1027
+ f"Manifest contains path traversal: {relpath!r}"
1028
+ )
1029
+ missing.append((entry, file_path, file_cache_path))
1030
+
1031
+ files, _ = _scan_entries(path, expected_paths)
1032
+ for relpath in files:
1033
+ if relpath in expected_paths:
1034
+ continue
1035
+ extra_path = path / relpath
1036
+ if extra_path.is_symlink():
1037
+ changed[relpath] = _get_symlink_cache_hash(extra_path, cache_dir)
1038
+ else:
1039
+ changed[relpath], _ = hash_file(extra_path, state_db)
1040
+
1041
+ if changed:
1042
+ # Replacing the directory deletes every file it holds; refuse if any is not cached.
1043
+ local_changes = sorted(
1044
+ rel for rel, content_hash in changed.items() if not _is_cached(cache_dir, content_hash)
1045
+ )
1046
+ if local_changes:
1047
+ return MissingRestore(
1048
+ restored=0, unavailable=0, local_changes=local_changes, replace=False
1049
+ )
1050
+ return MissingRestore(restored=0, unavailable=0, local_changes=[], replace=True)
1051
+
1052
+ unavailable = sum(1 for _, _, file_cache_path in missing if not file_cache_path.exists())
1053
+ if unavailable:
1054
+ return MissingRestore(restored=0, unavailable=unavailable, local_changes=[], replace=False)
1055
+
1056
+ for entry, file_path, file_cache_path in missing:
1057
+ file_path.parent.mkdir(parents=True, exist_ok=True)
1058
+ _checkout_with_fallback(
1059
+ file_path, file_cache_path, checkout_modes, executable=entry["isexec"]
1060
+ )
1061
+
1062
+ return MissingRestore(restored=len(missing), unavailable=0, local_changes=[], replace=False)
1063
+
1064
+
810
1065
  def protect(path: pathlib.Path) -> None:
811
1066
  """Make file read-only (mode 0o444)."""
812
1067
  os.chmod(path, 0o444)
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes