Cleanup

generator/cache.py

@@ -0,0 +1,82 @@
+"""Content identities and atomic storage shared by generator cache adapters."""
+from __future__ import annotations
+import ast
+import contextlib
+import functools
+import hashlib
+import json
+import os
+import pathlib
+import shutil
+import tempfile
+
+ROOT = pathlib.Path(__file__).resolve().parent
+
+@functools.lru_cache(maxsize=256)
+def imports(content):
+  modules = set()
+  for node in ast.walk(ast.parse(content)):
+    if isinstance(node, ast.Import):
+      modules.update(alias.name for alias in node.names if alias.name.startswith('generator.'))
+    elif isinstance(node, ast.ImportFrom) and node.module and node.module.startswith('generator'):
+      modules.add(node.module)
+      modules.update(node.module + '.' + alias.name for alias in node.names)
+  return tuple(sorted(modules))
+
+def source_identity(entrypoints):
+  """Conservatively include the transitive imports of explicit stage roots.
+
+  Read bytes on every call, so edits in a long-lived build process invalidate
+  both code content and newly introduced dependencies. Only parsing is cached.
+  """
+  pending, sources = list(entrypoints), {}
+  while pending:
+    module = pending.pop()
+    relative = module.split('.')[1:]
+    path = ROOT.joinpath(*relative).with_suffix('.py') if relative else ROOT / '__init__.py'
+    if not path.is_file():
+      path = ROOT.joinpath(*relative, '__init__.py')
+    if path in sources or not path.is_file():
+      continue
+    content = path.read_bytes()
+    sources[path] = hashlib.sha256(content).hexdigest()
+    pending.extend(imports(content))
+    pending.extend('.'.join(module.split('.')[:i]) for i in range(1, len(module.split('.'))))
+  return {path.relative_to(ROOT).as_posix(): sources[path] for path in sorted(sources)}
+
+def fingerprint(value):
+  return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
+
+def write_json(path, value):
+  path.parent.mkdir(parents=True, exist_ok=True)
+  temporary = None
+  try:
+    with tempfile.NamedTemporaryFile('w', dir=path.parent, delete=False) as output:
+      temporary = pathlib.Path(output.name)
+      json.dump(value, output, sort_keys=True)
+    os.replace(temporary, path)
+  finally:
+    if temporary is not None:
+      temporary.unlink(missing_ok=True)
+
+def has_symlink(path):
+  return any(parent.is_symlink() for parent in (path, *path.parents))
+
+@contextlib.contextmanager
+def stage_entry(entry, prefix):
+  if has_symlink(entry.parent):
+    raise ValueError('cache directories must not be symlinks')
+  entry.parent.mkdir(parents=True, exist_ok=True)
+  if entry.is_symlink() or entry.is_file():
+    entry.unlink(missing_ok=True)
+  entry.mkdir(exist_ok=True)
+  with tempfile.TemporaryDirectory(prefix=prefix, dir=entry.parent) as temporary:
+    yield pathlib.Path(temporary)
+
+def publish_files(entry, staged, names):
+  """Publish complete files in order, with the artifact manifest last."""
+  for name in names:
+    target = entry / name
+    if target.is_dir() and not target.is_symlink():
+      shutil.rmtree(target)
+    (staged / name).replace(target)