"""Read-only STRING species snapshot loader and PPI species resolver.""" from __future__ import annotations import hashlib import json from dataclasses import dataclass from functools import lru_cache from pathlib import Path from types import MappingProxyType from typing import Mapping from scripts.prepare_string_species_snapshot import ( PARSER_SCHEMA_VERSION, TreeNode, read_tree_nodes, ) VIRIDIPLANTAE_TAXON_ID = "33090" ARABIDOPSIS_TAXON_ID = "3702" PROJECT_SPECIES_ALIASES = { "at": ARABIDOPSIS_TAXON_ID, "ft": "62330", "os": "39947", } class StringSnapshotError(RuntimeError): def __init__(self, message: str, *, code: str) -> None: super().__init__(message) self.code = code @dataclass(frozen=True, slots=True) class StringSpeciesRecord: taxon_id: str string_type: str compact_name: str official_name: str domain: str @dataclass(frozen=True, slots=True) class StringSpeciesSnapshot: root: Path string_version: str snapshot_id: str nodes: Mapping[str, TreeNode] registry: Mapping[str, StringSpeciesRecord] name_to_taxa: Mapping[str, tuple[str, ...]] plant_string_taxa: tuple[str, ...] def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _normalize_name(value: str) -> str: return " ".join(str(value or "").split()).casefold() def _read_registry(path: Path) -> dict[str, StringSpeciesRecord]: records: dict[str, StringSpeciesRecord] = {} with path.open("r", encoding="utf-8", errors="strict") as handle: for line_number, raw in enumerate(handle, start=1): line = raw.rstrip("\r\n") if not line: continue fields = line.split("\t") if fields[0].lstrip("#").strip().lower() == "taxon_id": continue if raw.startswith("#"): continue if len(fields) < 5 or not fields[0].strip().isdigit(): raise ValueError(f"invalid STRING species row at line {line_number}") taxon_id = fields[0].strip() if taxon_id in records: raise ValueError(f"duplicate STRING registry taxon ID: {taxon_id}") records[taxon_id] = StringSpeciesRecord( taxon_id=taxon_id, string_type=fields[1].strip(), compact_name=fields[2].strip(), official_name=fields[3].strip(), domain=fields[4].strip(), ) if not records: raise ValueError("STRING species registry is empty") return records def _lineage_contains( nodes: Mapping[str, TreeNode], taxon_id: str, ancestor_taxon_id: str ) -> bool: current = taxon_id while True: if current == ancestor_taxon_id: return True node = nodes[current] if node.parent_taxon_id == current: return False current = node.parent_taxon_id def _validate_file(path: Path, metadata: object) -> None: if not path.is_file() or not isinstance(metadata, dict): raise StringSnapshotError( f"STRING snapshot file unavailable: {path}", code="string_registry_unavailable", ) expected_size = metadata.get("size") expected_sha256 = metadata.get("sha256") if ( not isinstance(expected_size, int) or not isinstance(expected_sha256, str) or path.stat().st_size != expected_size or _sha256(path) != expected_sha256 ): raise StringSnapshotError( f"STRING snapshot integrity check failed: {path}", code="string_registry_integrity_failed", ) def load_string_snapshot(root: Path) -> StringSpeciesSnapshot: root = Path(root).resolve() manifest_path = root / "source_manifest.json" if not manifest_path.is_file(): raise StringSnapshotError( f"STRING snapshot manifest unavailable: {manifest_path}", code="string_registry_unavailable", ) try: manifest = json.loads(manifest_path.read_text(encoding="utf-8")) snapshot_id = str(manifest["snapshot_id"]) string_version = str(manifest["string_version"]) files = manifest["files"] if not isinstance(files, dict): raise ValueError("files must be an object") species_metadata = files.get("species.v12.0.txt") tree_metadata = files.get("species.tree.v12.0.txt") if not isinstance(species_metadata, dict) or not isinstance(tree_metadata, dict): raise ValueError("snapshot file metadata must be objects") species_sha256 = str(species_metadata["sha256"]) tree_sha256 = str(tree_metadata["sha256"]) species_size = int(species_metadata["size"]) tree_size = int(tree_metadata["size"]) except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: raise StringSnapshotError( f"invalid STRING snapshot manifest: {manifest_path}: {exc}", code="string_registry_integrity_failed", ) from exc return _load_string_snapshot_cached( str(root), string_version, snapshot_id, species_sha256, tree_sha256, species_size, tree_size, ) @lru_cache(maxsize=4) def _load_string_snapshot_cached( root_text: str, string_version: str, snapshot_id: str, species_sha256: str, tree_sha256: str, species_size: int, tree_size: int, ) -> StringSpeciesSnapshot: root = Path(root_text) species_path = root / "species.v12.0.txt" tree_path = root / "species.tree.v12.0.txt" _validate_file( species_path, {"size": species_size, "sha256": species_sha256}, ) _validate_file( tree_path, {"size": tree_size, "sha256": tree_sha256}, ) expected_snapshot_id = hashlib.sha256( ( PARSER_SCHEMA_VERSION + "\n" + species_sha256 + "\n" + tree_sha256 ).encode("utf-8") ).hexdigest() if snapshot_id != expected_snapshot_id: raise StringSnapshotError( "STRING snapshot_id does not match its file hashes", code="string_registry_integrity_failed", ) try: registry = _read_registry(species_path) nodes = read_tree_nodes(tree_path) registry_taxa = set(registry) tree_string_taxa = { taxon_id for taxon_id, node in nodes.items() if node.is_string_species } if registry_taxa != tree_string_taxa: raise ValueError("STRING registry/tree membership mismatch") if VIRIDIPLANTAE_TAXON_ID not in nodes or ARABIDOPSIS_TAXON_ID not in registry: raise ValueError("required plant reference nodes are missing") except (OSError, UnicodeError, ValueError) as exc: raise StringSnapshotError( f"STRING snapshot parsing failed: {exc}", code="string_registry_integrity_failed", ) from exc names: dict[str, set[str]] = {} for taxon_id, node in nodes.items(): normalized = _normalize_name(node.taxon_name) if normalized: names.setdefault(normalized, set()).add(taxon_id) for taxon_id, record in registry.items(): for value in (record.compact_name, record.official_name): normalized = _normalize_name(value) if normalized: names.setdefault(normalized, set()).add(taxon_id) plant_string_taxa = tuple( sorted( ( taxon_id for taxon_id in registry if _lineage_contains(nodes, taxon_id, VIRIDIPLANTAE_TAXON_ID) ), key=int, ) ) return StringSpeciesSnapshot( root=root, string_version=string_version, snapshot_id=snapshot_id, nodes=MappingProxyType(nodes), registry=MappingProxyType(registry), name_to_taxa=MappingProxyType( { name: tuple(sorted(taxa, key=int)) for name, taxa in names.items() } ), plant_string_taxa=plant_string_taxa, ) def _unresolved_result( snapshot: StringSpeciesSnapshot, *, resolution_source: str, allow_model_fallback: bool, ) -> dict[str, object]: if allow_model_fallback: arabidopsis = snapshot.nodes[ARABIDOPSIS_TAXON_ID] return { "selection_mode": "model_fallback", "target_tree_taxon_id": "", "target_tree_name": "", "selected_string_taxon_id": ARABIDOPSIS_TAXON_ID, "selected_string_name": arabidopsis.taxon_name, "target_string_supported": False, "requires_user_confirmation": True, "resolution_source": resolution_source, "candidate_source": "explicit_model_fallback", "candidates": [], "issue_code": "ppi_model_fallback_confirmation_required", "string_version": snapshot.string_version, "snapshot_id": snapshot.snapshot_id, } return { "selection_mode": "unresolved", "target_tree_taxon_id": "", "target_tree_name": "", "selected_string_taxon_id": "", "selected_string_name": "", "target_string_supported": False, "requires_user_confirmation": False, "resolution_source": resolution_source, "candidate_source": "", "candidates": [], "issue_code": "target_string_tree_unresolved", "string_version": snapshot.string_version, "snapshot_id": snapshot.snapshot_id, } def _resolve_target_taxon( snapshot: StringSpeciesSnapshot, project_species_code: str, target_scientific_name: str, explicit_taxon_id: str, ) -> tuple[str, str]: explicit = str(explicit_taxon_id or "").strip() if explicit: return (explicit, "explicit_tree_taxon_id") if explicit in snapshot.nodes else ("", "explicit_tree_taxon_id") scientific_name = _normalize_name(target_scientific_name) if scientific_name: matches = snapshot.name_to_taxa.get(scientific_name, ()) return (matches[0], "scientific_name") if len(matches) == 1 else ("", "scientific_name") alias = str(project_species_code or "").strip().casefold() taxon_id = PROJECT_SPECIES_ALIASES.get(alias, "") return (taxon_id, "project_alias") if taxon_id in snapshot.nodes else ("", "project_alias") def _candidate_for( snapshot: StringSpeciesSnapshot, target_ancestors: Mapping[str, int], candidate_taxon_id: str, ) -> dict[str, object]: current = candidate_taxon_id candidate_edges = 0 while current not in target_ancestors: parent = snapshot.nodes[current].parent_taxon_id if parent == current: break current = parent candidate_edges += 1 lca_taxon_id = current target_edges = int(target_ancestors[lca_taxon_id]) node = snapshot.nodes[candidate_taxon_id] return { "selected_string_taxon_id": candidate_taxon_id, "selected_string_name": node.taxon_name, "lca_taxon_id": lca_taxon_id, "lca_name": snapshot.nodes[lca_taxon_id].taxon_name, "target_to_lca_edges": target_edges, "candidate_to_lca_edges": candidate_edges, "tree_distance": target_edges + candidate_edges, } def resolve_ppi_string_species( snapshot: StringSpeciesSnapshot, project_species_code: str, target_scientific_name: str = "", ppi_target_tree_taxon_id: str = "", top_n: int = 3, allow_model_fallback: bool = False, ) -> dict[str, object]: if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n < 1: raise ValueError("top_n must be a positive integer") target_taxon_id, resolution_source = _resolve_target_taxon( snapshot, project_species_code, target_scientific_name, ppi_target_tree_taxon_id, ) if not target_taxon_id: return _unresolved_result( snapshot, resolution_source=resolution_source, allow_model_fallback=allow_model_fallback, ) target_node = snapshot.nodes[target_taxon_id] if target_node.is_string_species and target_taxon_id in snapshot.registry: return { "selection_mode": "exact_string", "target_tree_taxon_id": target_taxon_id, "target_tree_name": target_node.taxon_name, "selected_string_taxon_id": target_taxon_id, "selected_string_name": target_node.taxon_name, "target_string_supported": True, "requires_user_confirmation": False, "resolution_source": resolution_source, "candidate_source": "string_species_registry", "candidates": [], "issue_code": "", "string_version": snapshot.string_version, "snapshot_id": snapshot.snapshot_id, } target_ancestors: dict[str, int] = {} current = target_taxon_id distance = 0 while current not in target_ancestors: target_ancestors[current] = distance parent = snapshot.nodes[current].parent_taxon_id if parent == current: break current = parent distance += 1 candidates = [ _candidate_for(snapshot, target_ancestors, candidate_taxon_id) for candidate_taxon_id in snapshot.plant_string_taxa ] candidates.sort( key=lambda row: ( int(row["tree_distance"]), int(row["selected_string_taxon_id"]), ) ) candidates = candidates[:top_n] if not candidates: return _unresolved_result( snapshot, resolution_source=resolution_source, allow_model_fallback=allow_model_fallback, ) return { "selection_mode": "tree_nearest", "target_tree_taxon_id": target_taxon_id, "target_tree_name": target_node.taxon_name, "selected_string_taxon_id": "", "selected_string_name": "", "target_string_supported": False, "requires_user_confirmation": True, "resolution_source": resolution_source, "candidate_source": "string_species_tree", "candidates": candidates, "issue_code": "ppi_reference_confirmation_required", "string_version": snapshot.string_version, "snapshot_id": snapshot.snapshot_id, }