data-annotations 2.1.2__py3-none-any.whl

data_annotations/provenance/git.py

@@ -0,0 +1,121 @@
+import subprocess
+import re
+from typing import Any
+from urllib.parse import urlsplit, urlunsplit
+
+
+_SCP_LIKE_REMOTE_RE = re.compile(r"^(?P<user>[^@]+)@(?P<host>[^:]+):(?P<path>.+)$")
+
+
+def _git_output(*args: str, subprocess_module=subprocess) -> str | None:
+    try:
+        return subprocess_module.check_output(
+            ["git", *args],
+            text=True,
+            stderr=subprocess_module.DEVNULL,
+        ).strip()
+    except Exception:
+        return None
+
+
+def _strip_git_suffix(path: str) -> str:
+    if path.endswith(".git"):
+        return path[:-4]
+    return path
+
+
+def _sanitize_remote_url(remote_url: str | None) -> str | None:
+    if not remote_url:
+        return None
+
+    scp_match = _SCP_LIKE_REMOTE_RE.match(remote_url)
+    if scp_match is not None:
+        host = scp_match.group("host")
+        path = _strip_git_suffix(scp_match.group("path").lstrip("/"))
+        return f"https://{host}/{path}"
+
+    parsed = urlsplit(remote_url)
+    hostname = parsed.hostname
+    if not hostname:
+        return remote_url
+
+    netloc = hostname
+    if parsed.port is not None:
+        netloc = f"{netloc}:{parsed.port}"
+
+    path = parsed.path
+    if parsed.scheme in {"http", "https"}:
+        return urlunsplit((parsed.scheme, netloc, _strip_git_suffix(path), "", ""))
+
+    if parsed.scheme in {"ssh", "git"} and path:
+        return urlunsplit(("https", netloc, _strip_git_suffix(path), "", ""))
+
+    return urlunsplit((parsed.scheme, netloc, path, parsed.query, parsed.fragment))
+
+
+def _remote_names(*, git_output_fn) -> list[str]:
+    remote_output = git_output_fn("remote")
+    if not remote_output:
+        return []
+    return [line.strip() for line in remote_output.splitlines() if line.strip()]
+
+
+def _upstream_remote_name(*, git_output_fn) -> str | None:
+    upstream = git_output_fn(
+        "rev-parse",
+        "--abbrev-ref",
+        "--symbolic-full-name",
+        "@{upstream}",
+    )
+    if not upstream or "/" not in upstream:
+        return None
+    return upstream.split("/", 1)[0]
+
+
+def _canonical_remote_name(*, git_output_fn) -> str | None:
+    upstream_remote = _upstream_remote_name(git_output_fn=git_output_fn)
+    if upstream_remote and git_output_fn("remote", "get-url", upstream_remote):
+        return upstream_remote
+
+    if git_output_fn("remote", "get-url", "origin"):
+        return "origin"
+
+    remotes = _remote_names(git_output_fn=git_output_fn)
+    if len(remotes) == 1 and git_output_fn("remote", "get-url", remotes[0]):
+        return remotes[0]
+
+    return None
+
+
+def capture_git_info(
+    *,
+    git_output_fn=None,
+    subprocess_module=None,
+) -> dict[str, Any]:
+    git_output_fn = _git_output if git_output_fn is None else git_output_fn
+    subprocess_module = subprocess if subprocess_module is None else subprocess_module
+    sha = git_output_fn("rev-parse", "HEAD")
+    branch = git_output_fn("rev-parse", "--abbrev-ref", "HEAD")
+    remote_name = _canonical_remote_name(git_output_fn=git_output_fn)
+    remote_url = _sanitize_remote_url(
+        git_output_fn("remote", "get-url", remote_name) if remote_name else None
+    )
+    try:
+        dirty = (
+            subprocess_module.call(
+                ["git", "diff", "--quiet"],
+                stdout=subprocess_module.DEVNULL,
+                stderr=subprocess_module.DEVNULL,
+            )
+            != 0
+        )
+    except Exception:
+        dirty = None
+
+    return {
+        "git_sha": sha,
+        "git_branch": branch,
+        "git_dirty": dirty,
+        "git_remote_name": remote_name,
+        "git_remote_url": remote_url,
+    }

data_annotations/provenance/models.py

@@ -0,0 +1,50 @@
+from datetime import datetime
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field
+
+ArtifactKind = Literal["plot", "model", "table", "dataset", "report", "other"]
+
+
+class ProducedFile(BaseModel):
+    path: str
+    kind: ArtifactKind = "other"
+    sha256: str | None = None
+
+
+class BaseProvenance(BaseModel):
+    created_at: datetime
+    capture_mode: Literal["runtime", "post_hoc"] = "runtime"
+    hostname: str
+    username: str
+    script: str | None = None
+    script_repo_path: str | None = None
+    function: str | None = None
+    command: list[str] = Field(default_factory=list)
+    params: dict[str, Any] = Field(default_factory=dict)
+    inputs: list[str] = Field(default_factory=list)
+    git_sha: str | None = None
+    git_branch: str | None = None
+    git_dirty: bool | None = None
+    git_remote_name: str | None = None
+    git_remote_url: str | None = None
+    slurm_job_id: str | None = None
+
+
+class DirectoryManifest(BaseProvenance):
+    output_dir: str
+    produced_files: list[ProducedFile] = Field(default_factory=list)
+
+
+class FileManifest(BaseProvenance):
+    artifact_path: str
+    artifact_kind: ArtifactKind = "other"
+    artifact_sha256: str | None = None
+
+
+class RecoveredSource(BaseModel):
+    checkout_path: str
+    script_path: str | None = None
+    git_sha: str
+    git_remote_url: str
+    script_repo_path: str | None = None

data_annotations/provenance/recovery.py

@@ -0,0 +1,473 @@
+import hashlib
+import json
+import os
+import re
+import subprocess
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+from .models import BaseProvenance, DirectoryManifest, FileManifest, RecoveredSource
+from .writers import sha256_file
+
+MatchStatus = Literal["match", "partial_match", "mismatch", "unverifiable"]
+_IGNORED_DIRECTORY_EXTRAS = frozenset({"manifest.json", "README.md", "schema.json"})
+_IGNORED_DIRECTORY_SUFFIXES = (".meta.json", ".README.md", ".prov.json", ".schema.json")
+
+
+def _is_ignored_directory_extra(path: Path) -> bool:
+    path_name = path.name
+    return path_name in _IGNORED_DIRECTORY_EXTRAS or path_name.endswith(
+        _IGNORED_DIRECTORY_SUFFIXES
+    )
+
+
+@dataclass(frozen=True)
+class ManifestMatch:
+    manifest_kind: Literal["file", "directory"]
+    status: MatchStatus
+    verified_entries: tuple[str, ...] = ()
+    missing_tracked_entries: tuple[str, ...] = ()
+    mismatched_tracked_entries: tuple[str, ...] = ()
+    extra_entries: tuple[str, ...] = ()
+    unverifiable_tracked_entries: tuple[str, ...] = ()
+
+
+def _file_annotation_manifest(payload: dict[str, Any]) -> FileManifest:
+    subject = payload["subject"]
+    provenance = payload["provenance"]
+    return FileManifest.model_validate(
+        {
+            **provenance,
+            "artifact_path": subject["path"],
+            "artifact_kind": subject.get("kind", "other"),
+            "artifact_sha256": subject.get("sha256"),
+        }
+    )
+
+
+def _directory_annotation_manifest(payload: dict[str, Any]) -> DirectoryManifest:
+    subject = payload["subject"]
+    provenance = payload["provenance"]
+    return DirectoryManifest.model_validate(
+        {
+            **provenance,
+            "output_dir": subject["path"],
+            "produced_files": subject.get("produced_files", []),
+        }
+    )
+
+
+def _normalize_allowed_value_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    normalized = dict(payload)
+    if "description" in normalized and "summary" not in normalized:
+        normalized["summary"] = normalized.pop("description")
+    return normalized
+
+
+def _normalize_field_definition_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    normalized = dict(payload)
+    if "description" in normalized and "summary" not in normalized:
+        normalized["summary"] = normalized.pop("description")
+    normalized["allowed_values"] = [
+        _normalize_allowed_value_payload(value)
+        for value in normalized.get("allowed_values", [])
+    ]
+    return normalized
+
+
+def _normalize_artifact_description_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    normalized = dict(payload)
+    if "description" in normalized and "summary" not in normalized:
+        normalized["summary"] = normalized.pop("description")
+    normalized["fields"] = [
+        _normalize_field_definition_payload(field)
+        for field in normalized.get("fields", [])
+    ]
+    return normalized
+
+
+def _normalize_description_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    normalized = dict(payload)
+    if "description" in normalized and "summary" not in normalized:
+        normalized["summary"] = normalized.pop("description")
+    if (
+        "description_updated_at" in normalized
+        and "description_updated_at" not in normalized
+    ):
+        normalized["description_updated_at"] = normalized.pop("description_updated_at")
+    normalized["fields"] = [
+        _normalize_field_definition_payload(field)
+        for field in normalized.get("fields", [])
+    ]
+    normalized["artifacts"] = [
+        _normalize_artifact_description_payload(artifact)
+        for artifact in normalized.get("artifacts", [])
+    ]
+    return normalized
+
+
+def _normalize_annotation_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    if payload.get("annotation_version") == "3":
+        return payload
+
+    normalized = dict(payload)
+    normalized["annotation_version"] = "3"
+    description = normalized.get("description")
+    if isinstance(description, dict):
+        normalized["description"] = _normalize_description_payload(description)
+    return normalized
+
+
+def _load_annotation_document(manifest: str | Path):
+    from data_annotations.annotations import (
+        DirectoryAnnotationDocument,
+        FileAnnotationDocument,
+    )
+
+    payload = json.loads(Path(manifest).read_text(encoding="utf-8"))
+    if (
+        payload.get("annotation_version") in {"2", "3"}
+        and isinstance(payload.get("subject"), dict)
+        and isinstance(payload.get("provenance"), dict)
+    ):
+        payload = _normalize_annotation_payload(payload)
+        if "produced_files" in payload["subject"]:
+            return DirectoryAnnotationDocument.model_validate(payload)
+        return FileAnnotationDocument.model_validate(payload)
+    raise ValueError("annotation payload is not a supported annotation document")
+
+
+def _parse_manifest_payload(payload: dict) -> FileManifest | DirectoryManifest:
+    if (
+        payload.get("annotation_version") in {"2", "3"}
+        and isinstance(payload.get("subject"), dict)
+        and isinstance(payload.get("provenance"), dict)
+    ):
+        subject = payload["subject"]
+        if "produced_files" in subject:
+            return _directory_annotation_manifest(payload)
+        if "path" in subject:
+            return _file_annotation_manifest(payload)
+    if "artifact_path" in payload:
+        return FileManifest.model_validate(payload)
+    if "output_dir" in payload:
+        return DirectoryManifest.model_validate(payload)
+    raise ValueError("annotation payload is not a supported provenance document")
+
+
+def _load_manifest(
+    manifest: str | Path | BaseProvenance,
+) -> FileManifest | DirectoryManifest:
+    if isinstance(manifest, FileManifest | DirectoryManifest):
+        return manifest
+    if isinstance(manifest, BaseProvenance):
+        raise ValueError("manifest must be a file or directory provenance manifest")
+
+    manifest_path = Path(manifest)
+    payload = json.loads(manifest_path.read_text(encoding="utf-8"))
+    return _parse_manifest_payload(payload)
+
+
+def _path_label(path: Path) -> str:
+    return path.as_posix()
+
+
+def _tracked_relative_path(
+    *,
+    produced_path: str | Path,
+    output_dir: str | Path,
+) -> Path | None:
+    produced_path = Path(produced_path)
+    output_dir = Path(output_dir)
+
+    if produced_path.is_absolute():
+        try:
+            return produced_path.relative_to(output_dir)
+        except ValueError:
+            return None
+
+    try:
+        return produced_path.relative_to(output_dir)
+    except ValueError:
+        return produced_path
+
+
+def _analyze_file_match(
+    artifact_path: Path,
+    manifest: FileManifest,
+) -> ManifestMatch:
+    entry_name = artifact_path.name
+    if manifest.artifact_sha256 is None:
+        return ManifestMatch(
+            manifest_kind="file",
+            status="unverifiable",
+            unverifiable_tracked_entries=(entry_name,),
+        )
+
+    if sha256_file(artifact_path) == manifest.artifact_sha256:
+        return ManifestMatch(
+            manifest_kind="file",
+            status="match",
+            verified_entries=(entry_name,),
+        )
+
+    return ManifestMatch(
+        manifest_kind="file",
+        status="mismatch",
+        mismatched_tracked_entries=(entry_name,),
+    )
+
+
+def _analyze_directory_match(
+    artifact_path: Path,
+    manifest: DirectoryManifest,
+) -> ManifestMatch:
+    verified_entries: list[str] = []
+    missing_tracked_entries: list[str] = []
+    mismatched_tracked_entries: list[str] = []
+    unverifiable_tracked_entries: list[str] = []
+    tracked_relative_paths: set[Path] = set()
+
+    for produced_file in manifest.produced_files:
+        relative_path = _tracked_relative_path(
+            produced_path=produced_file.path,
+            output_dir=manifest.output_dir,
+        )
+        if relative_path is None:
+            unverifiable_tracked_entries.append(str(produced_file.path))
+            continue
+
+        tracked_relative_paths.add(relative_path)
+        candidate_path = artifact_path / relative_path
+        relative_label = _path_label(relative_path)
+
+        if not candidate_path.exists() or not candidate_path.is_file():
+            missing_tracked_entries.append(relative_label)
+            continue
+
+        if produced_file.sha256 is None:
+            unverifiable_tracked_entries.append(relative_label)
+            continue
+
+        if sha256_file(candidate_path) == produced_file.sha256:
+            verified_entries.append(relative_label)
+        else:
+            mismatched_tracked_entries.append(relative_label)
+
+    candidate_relative_paths = {
+        path.relative_to(artifact_path)
+        for path in artifact_path.rglob("*")
+        if path.is_file()
+    }
+    extra_entries = sorted(
+        _path_label(path)
+        for path in candidate_relative_paths
+        if path not in tracked_relative_paths and not _is_ignored_directory_extra(path)
+    )
+
+    verified_entries = sorted(verified_entries)
+    missing_tracked_entries = sorted(missing_tracked_entries)
+    mismatched_tracked_entries = sorted(mismatched_tracked_entries)
+    unverifiable_tracked_entries = sorted(unverifiable_tracked_entries)
+
+    if mismatched_tracked_entries:
+        status: MatchStatus = "mismatch"
+    elif verified_entries:
+        if missing_tracked_entries or extra_entries or unverifiable_tracked_entries:
+            status = "partial_match"
+        else:
+            status = "match"
+    else:
+        status = "unverifiable"
+
+    return ManifestMatch(
+        manifest_kind="directory",
+        status=status,
+        verified_entries=tuple(verified_entries),
+        missing_tracked_entries=tuple(missing_tracked_entries),
+        mismatched_tracked_entries=tuple(mismatched_tracked_entries),
+        extra_entries=tuple(extra_entries),
+        unverifiable_tracked_entries=tuple(unverifiable_tracked_entries),
+    )
+
+
+def _analyze_artifact_match(
+    artifact_path: str | Path,
+    manifest: str | Path | BaseProvenance,
+) -> ManifestMatch:
+    loaded_manifest = _load_manifest(manifest)
+    artifact_path = Path(artifact_path).expanduser().resolve()
+
+    if isinstance(loaded_manifest, FileManifest):
+        if artifact_path.is_dir():
+            raise ValueError("file manifest matching requires a file target")
+        return _analyze_file_match(artifact_path, loaded_manifest)
+
+    if not artifact_path.is_dir():
+        raise ValueError("directory manifest matching requires a directory target")
+    return _analyze_directory_match(artifact_path, loaded_manifest)
+
+
+def artifact_matches_manifest(
+    artifact_path: str | Path,
+    manifest: str | Path | FileManifest | DirectoryManifest,
+) -> bool:
+    match = _analyze_artifact_match(artifact_path, manifest)
+    return match.status in {"match", "partial_match"}
+
+
+def _safe_remote_slug(remote_url: str) -> str:
+    normalized = remote_url.removesuffix(".git").rstrip("/").replace("\\", "/")
+    slug = re.sub(r"[^A-Za-z0-9._-]+", "-", normalized).strip("-") or "repo"
+    digest = hashlib.sha1(remote_url.encode("utf-8")).hexdigest()[:8]
+    return f"{slug[-64:]}-{digest}"
+
+
+def _checkout_cache_root(
+    *,
+    env: dict[str, str] | os._Environ[str] | None = None,
+    home: Path | None = None,
+    platform: str | None = None,
+) -> Path:
+    env = os.environ if env is None else env
+    home = Path.home() if home is None else home
+    platform = sys.platform if platform is None else platform
+
+    if platform == "win32":
+        base_dir = Path(env.get("LOCALAPPDATA", home / "AppData" / "Local"))
+    elif platform == "darwin":
+        base_dir = home / "Library" / "Caches"
+    else:
+        base_dir = Path(env.get("XDG_CACHE_HOME", home / ".cache"))
+
+    return base_dir.expanduser().resolve() / "data-annotations" / "checkouts"
+
+
+def _default_checkout_path(
+    *,
+    git_remote_url: str,
+    git_sha: str,
+    env: dict[str, str] | os._Environ[str] | None = None,
+    home: Path | None = None,
+    platform: str | None = None,
+) -> Path:
+    return (
+        _checkout_cache_root(env=env, home=home, platform=platform)
+        / _safe_remote_slug(git_remote_url)
+        / git_sha
+    )
+
+
+def _git_output(*args: str) -> str | None:
+    try:
+        return subprocess.check_output(
+            ["git", *args],
+            text=True,
+            stderr=subprocess.DEVNULL,
+        ).strip()
+    except Exception:
+        return None
+
+
+def _checkout_matches_recorded_state(
+    checkout_path: Path,
+    *,
+    git_sha: str,
+    git_remote_url: str,
+) -> bool:
+    if not checkout_path.is_dir():
+        return False
+
+    head_sha = _git_output("-C", str(checkout_path), "rev-parse", "HEAD")
+    remote_url = _git_output(
+        "-C", str(checkout_path), "config", "--get", "remote.origin.url"
+    )
+    return head_sha == git_sha and remote_url == git_remote_url
+
+
+def _prepare_checkout_path(
+    checkout_path: Path,
+    *,
+    git_sha: str,
+    git_remote_url: str,
+) -> tuple[Path, bool]:
+    checkout_path = checkout_path.expanduser().resolve()
+
+    if checkout_path.exists():
+        if _checkout_matches_recorded_state(
+            checkout_path,
+            git_sha=git_sha,
+            git_remote_url=git_remote_url,
+        ):
+            return checkout_path, True
+
+        if checkout_path.is_dir() and not any(checkout_path.iterdir()):
+            return checkout_path, False
+
+        raise ValueError(
+            f"checkout destination already exists and does not match the recorded repository state: {checkout_path}"
+        )
+
+    checkout_path.parent.mkdir(parents=True, exist_ok=True)
+    return checkout_path, False
+
+
+def checkout_manifest_source(
+    manifest: str | Path | BaseProvenance,
+    *,
+    destination_dir: str | Path | None = None,
+) -> RecoveredSource:
+    loaded_manifest = _load_manifest(manifest)
+    if loaded_manifest.git_remote_url is None:
+        raise ValueError("manifest does not record git_remote_url")
+    if loaded_manifest.git_sha is None:
+        raise ValueError("manifest does not record git_sha")
+
+    checkout_path = (
+        _default_checkout_path(
+            git_remote_url=loaded_manifest.git_remote_url,
+            git_sha=loaded_manifest.git_sha,
+        )
+        if destination_dir is None
+        else Path(destination_dir).expanduser().resolve()
+    )
+    checkout_path, reused = _prepare_checkout_path(
+        checkout_path,
+        git_sha=loaded_manifest.git_sha,
+        git_remote_url=loaded_manifest.git_remote_url,
+    )
+
+    if not reused:
+        subprocess.check_call(
+            ["git", "clone", loaded_manifest.git_remote_url, str(checkout_path)],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        subprocess.check_call(
+            [
+                "git",
+                "-C",
+                str(checkout_path),
+                "checkout",
+                "--detach",
+                loaded_manifest.git_sha,
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+
+    script_path: str | None = None
+    if loaded_manifest.script_repo_path is not None:
+        candidate = checkout_path / loaded_manifest.script_repo_path
+        if candidate.exists():
+            script_path = str(candidate.resolve())
+
+    return RecoveredSource(
+        checkout_path=str(checkout_path.resolve()),
+        script_path=script_path,
+        git_sha=loaded_manifest.git_sha,
+        git_remote_url=loaded_manifest.git_remote_url,
+        script_repo_path=loaded_manifest.script_repo_path,
+    )