API Reference¶

@dataclass
class Alignment(Asset):
    """BAM/CRAM alignment file asset.

    Carries reference_cid and r1_cid for provenance (PROV entity derivation).
    """

    _asset_key: ClassVar[str] = "alignment"
    sample_id: str = ""
    format: str = ""
    sorted: str = ""
    duplicates_marked: bool = False
    bqsr_applied: bool = False
    tool: str = ""
    reference_cid: str = ""
    r1_cid: str = ""

@dataclass
class AlignmentIndex(Asset):
    """BAI/CRAI alignment index file asset.

    Carries alignment_cid linking to the Alignment it indexes.
    """

    _asset_key: ClassVar[str] = "alignment_index"
    sample_id: str = ""
    alignment_cid: str = ""

@dataclass
class BQSRReport(Asset):
    """BQSR recalibration table produced by GATK BaseRecalibrator.

    Carries alignment_cid linking to the Alignment it was produced from.
    """

    _asset_key: ClassVar[str] = "bqsr_report"
    sample_id: str = ""
    tool: str = ""
    alignment_cid: str = ""

@dataclass
class DuplicateMetrics(Asset):
    """Duplicate metrics text file produced by GATK MarkDuplicates.

    Carries alignment_cid linking to the Alignment it was produced from.
    """

    _asset_key: ClassVar[str] = "duplicate_metrics"
    sample_id: str = ""
    tool: str = ""
    alignment_cid: str = ""

@dataclass
class Asset:
    """Base class for all typed file assets in Stargazer.

    Attributes:
        cid: Content identifier (CID) for the stored file
        path: Local filesystem path (set after download or upload)
        keyvalues: Free-form metadata for *bare* Asset instances — the
            catchall for records whose asset key has no registered class
            in this process. Serialized verbatim (the ``asset`` key lives
            inside the dict). Typed subclasses ignore this field entirely;
            they serialize their declared fields instead.

    Subclasses declare typed fields as normal dataclass attributes:

        @dataclass
        class Alignment(Asset):
            _asset_key: ClassVar[str] = "alignment"
            sample_id: str = ""
            duplicates_marked: bool = False

    Fields are plain Python attributes. ``to_keyvalues()`` serializes them to
    ``dict[str, str]`` at storage boundaries; ``from_keyvalues()`` reconstructs
    from storage. ``str`` fields pass through directly; all other types use
    ``json.dumps`` / ``json.loads``.
    """

    _registry: ClassVar[dict[str, type["Asset"]]] = {}
    _asset_key: ClassVar[str] = ""

    cid: str = ""
    path: Path | None = None
    keyvalues: dict[str, str] = field(default_factory=dict)

    def __init_subclass__(cls, **kwargs):
        """Register subclass in the asset registry."""
        super().__init_subclass__(**kwargs)
        ak = cls.__dict__.get("_asset_key", "")
        if ak:
            Asset._registry[ak] = cls

    def __setattr__(self, name: str, value: Any) -> None:
        """Enforce declared fields on typed subclasses; pass through on base Asset."""
        if self._asset_key and not name.startswith("_") and name not in _BASE_FIELDS:
            allowed = {f.name for f in dataclasses.fields(type(self))} - _BASE_FIELDS
            if name not in allowed:
                raise AttributeError(
                    f"{type(self).__name__} has no field '{name}'. "
                    f"Allowed: {sorted(allowed)}"
                )
        super().__setattr__(name, value)

    def to_keyvalues(self) -> dict[str, str]:
        """Serialize to storage format.

        str fields pass through as-is; all other types are serialized with
        json.dumps. Bare Asset instances return a copy of their keyvalues
        dict verbatim — a copy so callers (e.g. ``_owner`` stamping at the
        storage layer) can mutate the result without touching the Asset.
        """
        if not self._asset_key:
            return dict(self.keyvalues)
        hints = get_type_hints(type(self))
        result: dict[str, str] = {"asset": self._asset_key}
        for f in dataclasses.fields(self):
            if f.name in _BASE_FIELDS:
                continue
            val = getattr(self, f.name)
            result[f.name] = val if hints.get(f.name) is str else json.dumps(val)
        return result

    @classmethod
    def from_keyvalues(
        cls, kv: dict[str, str], cid: str = "", path: Path | None = None
    ) -> "Asset":
        """Reconstruct from a storage keyvalues dict.

        str fields are assigned directly; all other types are deserialized
        with json.loads (raising ``json.JSONDecodeError`` on values that
        don't parse — callers that must tolerate malformed records catch it,
        see ``specialize()``). Bare Asset receives the keyvalues verbatim.
        """
        if not cls._asset_key:
            return cls(cid=cid, path=path, keyvalues=dict(kv))
        hints = get_type_hints(cls)
        kwargs = {}
        for f in dataclasses.fields(cls):
            if f.name in _BASE_FIELDS:
                continue
            if f.name in kv:
                kwargs[f.name] = (
                    kv[f.name] if hints.get(f.name) is str else json.loads(kv[f.name])
                )
        return cls(cid=cid, path=path, **kwargs)

    def to_dict(self) -> dict:
        """Serialize to a JSON-friendly dict."""
        return {
            "cid": self.cid,
            "path": str(self.path) if self.path else None,
            "keyvalues": self.to_keyvalues(),
        }

    @classmethod
    def from_dict(cls, data: dict) -> Self:
        """Reconstruct from a serialized dict."""
        return cls.from_keyvalues(
            data.get("keyvalues", {}),
            cid=data.get("cid", ""),
            path=Path(data["path"]) if data.get("path") else None,
        )

    async def fetch(self) -> None:
        """Download this asset and all its companions from storage.

        Downloads the asset itself, then queries storage for any assets linked
        via ``{_asset_key}_cid`` to auto-download companions (e.g. indices,
        mate reads).
        """
        import stargazer.utils.local_storage as _storage

        await _storage.default_client.download(self)

        if self._asset_key and self.cid:
            companions = await assemble(**{f"{self._asset_key}_cid": self.cid})
            for a in companions:
                await _storage.default_client.download(a)

    async def update(self, path: Path, **kwargs) -> None:
        """Upload file and set cid. Shared by all asset types."""
        from stargazer.utils.local_storage import default_client

        for key, value in kwargs.items():
            if value is not None:
                setattr(self, key, value)
        self.path = path
        await default_client.upload(self)

@dataclass
class R1(Asset):
    """R1 (forward) FASTQ read file asset.

    Carries mate_cid pointing to the paired R2 asset's CID (None for single-end).
    """

    _asset_key: ClassVar[str] = "r1"
    sample_id: str = ""
    mate_cid: str = ""
    sequencing_platform: str = ""

@dataclass
class R2(Asset):
    """R2 (reverse) FASTQ read file asset.

    Carries mate_cid pointing to the paired R1 asset's CID (None for single-end).
    """

    _asset_key: ClassVar[str] = "r2"
    sample_id: str = ""
    mate_cid: str = ""
    sequencing_platform: str = ""

@dataclass
class AlignerIndex(Asset):
    """Aligner index file asset (one file per index file for multi-file indices)."""

    _asset_key: ClassVar[str] = "aligner_index"
    build: str = ""
    aligner: str = ""
    reference_cid: str = ""

@dataclass
class Reference(Asset):
    """Reference FASTA file asset."""

    _asset_key: ClassVar[str] = "reference"
    build: str = ""

    @property
    def contigs(self) -> list[str]:
        """Read contig names from the companion .fai index.

        Requires fetch() to have been called first so the ReferenceIndex
        companion is downloaded alongside this reference.
        """
        if self.path is None:
            raise ValueError("Reference has no local path — call fetch() first")
        fai_path = self.path.parent / (self.path.name + ".fai")
        if not fai_path.exists():
            raise FileNotFoundError(
                f"Reference index not found at {fai_path}. "
                f"Run samtools_faidx first, then fetch() to download companions."
            )
        contigs = []
        with open(fai_path) as f:
            for line in f:
                name = line.split("\t", 1)[0].strip()
                if name:
                    contigs.append(name)
        return contigs

@dataclass
class ReferenceIndex(Asset):
    """FASTA index (.fai) file asset.

    Carries reference_cid linking back to the Reference it was built from.
    """

    _asset_key: ClassVar[str] = "reference_index"
    build: str = ""
    tool: str = ""
    reference_cid: str = ""

@dataclass
class SequenceDict(Asset):
    """Sequence dictionary (.dict) file asset."""

    _asset_key: ClassVar[str] = "sequence_dict"
    build: str = ""
    tool: str = ""
    reference_cid: str = ""

@dataclass
class AnnData(Asset):
    """AnnData (.h5ad) file asset for single-cell RNA-seq data.

    Tracks pipeline stage, cell/gene counts, and provenance through
    the scRNA-seq processing steps.
    """

    _asset_key: ClassVar[str] = "anndata"
    sample_id: str = ""
    n_obs: int = 0
    n_vars: int = 0
    stage: str = ""
    organism: str = ""
    source_cid: str = ""

@dataclass
class KnownSites(Asset):
    """Known variant sites VCF used for BQSR.

    Standalone asset — carries build and source fields, no container needed.
    """

    _asset_key: ClassVar[str] = "known_sites"
    build: str = ""
    resource_name: str = ""
    known: str = "false"
    training: str = "false"
    truth: str = "false"
    prior: str = "10"
    vqsr_mode: str = ""

@dataclass
class KnownSitesIndex(Asset):
    """VCF index (.idx) file for a KnownSites asset.

    Carries known_sites_cid linking to the KnownSites VCF it indexes.
    Fetched automatically alongside the VCF via Asset.fetch().
    """

    _asset_key: ClassVar[str] = "known_sites_index"
    known_sites_cid: str = ""

@dataclass
class VQSRModel(Asset):
    """VQSR recalibration model (.recal file + tranches path).

    Produced by VariantRecalibrator. The recal file is the primary path;
    the companion tranches file path is stored in keyvalues["tranches_path"].
    """

    _asset_key: ClassVar[str] = "vqsr_model"
    sample_id: str = ""
    mode: str = "SNP"
    tranches_path: str = ""
    build: str = ""
    variants_cid: str = ""

@dataclass
class Variants(Asset):
    """VCF/GVCF variant call file asset."""

    _asset_key: ClassVar[str] = "variants"
    sample_id: str = ""
    caller: str = ""
    variant_type: str = ""
    build: str = ""
    vqsr_mode: str = ""
    sample_count: int = 0
    source_samples: list = None

@dataclass
class VariantsIndex(Asset):
    """VCF index (.tbi) file asset.

    Carries variants_cid linking to the Variants file it indexes.
    """

    _asset_key: ClassVar[str] = "variants_index"
    sample_id: str = ""
    variants_cid: str = ""

@app.class_definition
@dataclass
class SampleSheet(Asset):
    """CSV of per-sample metadata for a cohort."""

    _asset_key: ClassVar[str] = "sample_sheet"
    cohort_id: str = ""
    n_samples: int = 0
    organism: str = ""

@app.class_definition
@dataclass
class CohortSummary(Asset):
    """JSON summary of a cohort: per-organism counts and totals."""

    _asset_key: ClassVar[str] = "cohort_summary"
    cohort_id: str = ""
    n_samples: int = 0
    n_organisms: int = 0
    sample_sheet_cid: str = ""

@dataclass
class TaskInfo:
    """Complete metadata about a registered task."""

    name: str
    category: str  # "task" or "workflow"
    description: str
    params: list[TaskParam]
    outputs: list[TaskOutput]
    task_obj: Any  # The Flyte task object