common/artifacts.py

"""Artifact provenance stamping and the generator/artifact registry."""
import hashlib
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple

from .paths import REPO


@dataclass(frozen=True)
class ArtifactSpec:
    """Owning script, source pool, and required top-level payload keys."""

    generator: Optional[str]
    pool: Optional[str]
    payload_keys: Tuple[str, ...]


REGISTRY = {
    'rules.json': ArtifactSpec(
        'choose.py', 'TRAIN2017', ('candidates', 'sizes', 'rules')),
    'eval.json': ArtifactSpec('verify.py', 'VAL5000', ('rules',)),
    'synth.json': ArtifactSpec(
        'synth.py', None, ('tool', 'target', 'variants')),
}


def sha256_of(path) -> str:
    """Content hash of a file."""
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


def provenance(generator: str, classifier=None, pool_info: Optional[dict] = None,
               **extra) -> dict:
    """Assemble a provenance block from repository-recoverable fields only."""
    block = {'generator': generator}
    if classifier is not None:
        path = Path(classifier)
        block['classifier'] = str(path.resolve().relative_to(REPO)).replace('\\', '/')
        block['classifier_sha256'] = sha256_of(path)
    if pool_info:
        block.update(pool_info)
    block.update(extra)
    return block


def write_artifact(path, payload: dict, *, generator: str, classifier=None,
                   pool_info: Optional[dict] = None, **extra):
    """Write `payload` beneath a provenance block and return the document."""
    doc = {'provenance': provenance(generator, classifier, pool_info, **extra)}
    doc.update(payload)
    Path(path).write_text(json.dumps(doc, indent=2) + '\n', encoding='utf-8')
    return doc


def read_artifact(path) -> dict:
    return json.loads(Path(path).read_text(encoding='utf-8'))