Source code for matmelt.core.stages.extract

"""The extract stage: one screened-in structure -> OPTIMADE-shaped record (3.4).

The thin ``Stage`` half of the pair; the logic lives entirely in
``extract_runner.py`` -- which, unlike the other runners, is **stdlib only**:
extraction is bookkeeping, not science, so it runs in any python anywhere
with nothing to install (there is no exit-3 path).

1:1 and the funnel's terminal stage: the runner reads the item's
``relaxed.cif`` (parsed by its own minimal CIF reader -- geometry is
best-effort, a parse failure degrades to a geometry-less record, never a
crash), ``screen_result.json`` (must say ``RELAXED``: extract serves only
screened-in materials), the optional ``structure.json`` provenance from
the structure stage and the optional ``calc_result.json`` from the
calculate stage (folded under ``_mcli.calculate`` -- healing-attempt
history included), and writes ``record.json`` -- an OPTIMADE
``structures``-shaped document ready for a database or an OPTIMADE-serving
layer, with the non-standard extras namespaced under ``_mcli``.

``validate_inputs`` refuses clearly when ``screen_result.json`` does not
parse to a ``RELAXED`` verdict -- someone starting mid-pipeline on a
screened-OUT (or hand-made) item is told exactly why, before any job runs.
"""

from __future__ import annotations

import json
from pathlib import Path
from typing import TYPE_CHECKING

from matmelt.core.jobspec import JobSpec
from matmelt.core.stages.base import (
    Stage,
    StageError,
    StageInputError,
    StageInputSpec,
    StageOutputSpec,
    StageVerdict,
)
from matmelt.core.stages.registry import register

if TYPE_CHECKING:
    from matmelt.core.jobspec import JobResult
    from matmelt.core.stages.base import WorkItem

#: The self-contained (stdlib-only) script shipped with every extract job.
RUNNER_NAME = "extract_runner.py"

#: Inputs from the screen stage -- declared as input artifacts so the Phase-4
#: staging transport knows to push them (in-place locally; staged as a no-op).
RELAXED_NAME = "relaxed.cif"
SCREEN_NAME = "screen_result.json"

#: Optional provenance from the structure stage (space group, lattice type);
#: folded into the record when present, never required.
META_NAME = "structure.json"

#: Optional result from the calculate stage (DFTB+ WorkChain outcome with
#: the full healing-attempt history); folded under ``_mcli.calculate`` when
#: present, never required -- extract stays startable without calculate.
CALC_NAME = "calc_result.json"

#: The OPTIMADE-shaped document this stage produces.
RECORD_NAME = "record.json"

#: How many characters of a failed runner's stderr to surface in the verdict.
_REASON_TAIL = 300




class ExtractStage(Stage):
    """Assemble one screened-in material's OPTIMADE-shaped ``record.json``.

    Parameters
    ----------
    python
        Interpreter to invoke on the backend this stage is wired to. The
        runner needs nothing beyond the standard library, so any python
        works -- the option exists for backend-environment symmetry with the
        other stages, never because a dependency demands it.

    Raises
    ------
    StageError
        When ``python`` is not a non-empty string.
    """

    # No runtime_requirements: stdlib only, by contract -- the base default
    # () IS the declaration, and the preflight reports it as always-ready.

    def __init__(self, python: str = "python3") -> None:
        if not isinstance(python, str) or not python.strip():
            raise StageError(f"extract: python must be a non-empty string, got {python!r}")
        self.runner_python = python

    @property
    def name(self) -> str:
        return "extract"

    @property
    def requires(self) -> StageInputSpec:
        return StageInputSpec(
            files=(RELAXED_NAME, SCREEN_NAME),
            description="the relaxed structure + RELAXED screening verdict from the screen stage",
        )

    @property
    def produces(self) -> StageOutputSpec:
        return StageOutputSpec(
            files=(RECORD_NAME,),
            description="OPTIMADE structures-shaped record for the results database",
        )



    def validate_inputs(self, item: WorkItem) -> None:
        """Refuse clearly unless the screening verdict actually says RELAXED.

        Extends the base file-presence check by reading the verdict, which is
        the mid-pipeline-start guard: extract serves only screened-in
        materials.

        Parameters
        ----------
        item
            The material about to be extracted.

        Raises
        ------
        StageInputError
            When a required file is missing or empty, when the screening
            verdict is unreadable or not valid JSON, or when it names any
            category other than ``RELAXED``.
        """
        super().validate_inputs(item)
        path = item.dir / SCREEN_NAME
        try:
            verdict = json.loads(path.read_text(encoding="utf-8"))
        except OSError as exc:
            raise StageInputError(
                f"stage 'extract' cannot run on item {item.id!r}: {SCREEN_NAME} is "
                f"unreadable: {exc}"
            ) from exc
        except ValueError as exc:
            raise StageInputError(
                f"stage 'extract' cannot run on item {item.id!r}: {SCREEN_NAME} is "
                f"not valid JSON: {exc}"
            ) from exc
        category = verdict.get("category") if isinstance(verdict, dict) else None
        if category != "RELAXED":
            raise StageInputError(
                f"stage 'extract' cannot run on item {item.id!r}: {SCREEN_NAME} says "
                f"category {category!r}, but extract serves only RELAXED (screened-in) "
                f"materials"
            )




    def build_jobspec(self, item: WorkItem) -> JobSpec:
        runner = Path(__file__).with_name(RUNNER_NAME)
        input_artifacts = {
            RUNNER_NAME: str(runner),
            RELAXED_NAME: str(item.dir / RELAXED_NAME),
            SCREEN_NAME: str(item.dir / SCREEN_NAME),
        }
        # The structure stage's provenance file and the calculate stage's
        # result are optional enrichments: an item that never ran those
        # stages legitimately lacks them. Declare each only when it exists
        # -- a declared artifact must be stageable.
        if (item.dir / META_NAME).is_file():
            input_artifacts[META_NAME] = str(item.dir / META_NAME)
        if (item.dir / CALC_NAME).is_file():
            input_artifacts[CALC_NAME] = str(item.dir / CALC_NAME)
        return JobSpec(
            command=[self.runner_python, RUNNER_NAME, "--id", item.id],
            input_artifacts=input_artifacts,
            output_artifacts=[RECORD_NAME],
            cwd=str(item.dir),
            scratch_inputs=frozenset({RUNNER_NAME}),
        )




    def interpret_result(self, item: WorkItem, result: JobResult) -> StageVerdict:
        """Judge a finished extract job.

        Parameters
        ----------
        item
            The material whose job finished.
        result
            What the backend reported.

        Returns
        -------
        StageVerdict
            ``ok`` on exit 0, meaning the record was written. Anything else is
            a fault carrying the runner's stderr tail, and **never terminal**:
            with a stdlib-only runner a failure is a broken input file or an
            environment accident, and the item already carries a RELAXED
            verdict -- science never rejects here.
        """
        if result.exit_code == 0:
            return StageVerdict(ok=True)
        tail = (result.stderr or result.stdout).strip()[-_REASON_TAIL:]
        code = "unknown" if result.exit_code is None else str(result.exit_code)
        return StageVerdict(ok=False, reason=f"extract runner exited {code}: {tail}")




register("extract", ExtractStage)