268 lines
9.1 KiB
Python
268 lines
9.1 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Download and validate the fixed STRING v12.0 species snapshot."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import hashlib
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import tempfile
|
||
|
|
import urllib.request
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
STRING_VERSION = "12.0"
|
||
|
|
PARSER_SCHEMA_VERSION = "1"
|
||
|
|
SOURCE_URLS = {
|
||
|
|
"species.v12.0.txt": "https://stringdb-downloads.org/download/species.v12.0.txt",
|
||
|
|
"species.tree.v12.0.txt": "https://stringdb-downloads.org/download/species.tree.v12.0.txt",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True, slots=True)
|
||
|
|
class TreeNode:
|
||
|
|
taxon_id: str
|
||
|
|
parent_taxon_id: str
|
||
|
|
taxon_name: str
|
||
|
|
is_string_species: bool
|
||
|
|
|
||
|
|
|
||
|
|
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 _parse_tf(value: str, *, line_number: int) -> bool:
|
||
|
|
if value == "t":
|
||
|
|
return True
|
||
|
|
if value == "f":
|
||
|
|
return False
|
||
|
|
raise ValueError(
|
||
|
|
f"invalid is_STRING_species value at line {line_number}: {value!r}"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def read_registry_taxa(path: Path) -> set[str]:
|
||
|
|
taxa: set[str] = set()
|
||
|
|
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")
|
||
|
|
first = fields[0].lstrip("#").strip().lower()
|
||
|
|
if first == "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 taxa:
|
||
|
|
raise ValueError(f"duplicate STRING registry taxon ID: {taxon_id}")
|
||
|
|
taxa.add(taxon_id)
|
||
|
|
if not taxa:
|
||
|
|
raise ValueError("STRING species registry is empty")
|
||
|
|
return taxa
|
||
|
|
|
||
|
|
|
||
|
|
def read_tree_nodes(path: Path) -> dict[str, TreeNode]:
|
||
|
|
nodes: dict[str, TreeNode] = {}
|
||
|
|
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")
|
||
|
|
first = fields[0].lstrip("#").strip().lower()
|
||
|
|
if first == "taxon_id":
|
||
|
|
continue
|
||
|
|
if raw.startswith("#"):
|
||
|
|
continue
|
||
|
|
if len(fields) != 4:
|
||
|
|
raise ValueError(f"invalid STRING tree row at line {line_number}")
|
||
|
|
taxon_id, parent_taxon_id, taxon_name, raw_is_string = (
|
||
|
|
field.strip() for field in fields
|
||
|
|
)
|
||
|
|
if not taxon_id.isdigit() or not parent_taxon_id.isdigit():
|
||
|
|
raise ValueError(f"invalid STRING tree taxon ID at line {line_number}")
|
||
|
|
if taxon_id in nodes:
|
||
|
|
raise ValueError(f"duplicate STRING tree taxon ID: {taxon_id}")
|
||
|
|
nodes[taxon_id] = TreeNode(
|
||
|
|
taxon_id=taxon_id,
|
||
|
|
parent_taxon_id=parent_taxon_id,
|
||
|
|
taxon_name=taxon_name,
|
||
|
|
is_string_species=_parse_tf(raw_is_string, line_number=line_number),
|
||
|
|
)
|
||
|
|
if not nodes:
|
||
|
|
raise ValueError("STRING species tree is empty")
|
||
|
|
return nodes
|
||
|
|
|
||
|
|
|
||
|
|
def _validate_parent_graph(nodes: dict[str, TreeNode]) -> None:
|
||
|
|
unresolved = sorted(
|
||
|
|
{
|
||
|
|
node.parent_taxon_id
|
||
|
|
for node in nodes.values()
|
||
|
|
if node.parent_taxon_id not in nodes
|
||
|
|
},
|
||
|
|
key=int,
|
||
|
|
)
|
||
|
|
if unresolved:
|
||
|
|
raise ValueError(f"unresolved parent taxon IDs: {unresolved[:10]}")
|
||
|
|
|
||
|
|
completed: set[str] = set()
|
||
|
|
for start in nodes:
|
||
|
|
if start in completed:
|
||
|
|
continue
|
||
|
|
path: list[str] = []
|
||
|
|
positions: dict[str, int] = {}
|
||
|
|
current = start
|
||
|
|
while current not in completed:
|
||
|
|
if current in positions:
|
||
|
|
cycle = path[positions[current] :]
|
||
|
|
if len(cycle) == 1 and cycle[0] == "1":
|
||
|
|
break
|
||
|
|
raise ValueError(f"non-root cycle in STRING species tree: {cycle[:10]}")
|
||
|
|
positions[current] = len(path)
|
||
|
|
path.append(current)
|
||
|
|
parent = nodes[current].parent_taxon_id
|
||
|
|
if parent == current:
|
||
|
|
if current != "1":
|
||
|
|
raise ValueError(
|
||
|
|
f"non-root cycle in STRING species tree: {[current]}"
|
||
|
|
)
|
||
|
|
break
|
||
|
|
current = parent
|
||
|
|
completed.update(path)
|
||
|
|
|
||
|
|
|
||
|
|
def validate_snapshot(species_path: Path, tree_path: Path) -> dict[str, object]:
|
||
|
|
registry_taxa = read_registry_taxa(species_path)
|
||
|
|
nodes = read_tree_nodes(tree_path)
|
||
|
|
_validate_parent_graph(nodes)
|
||
|
|
|
||
|
|
for required in {"33090", "3702"}:
|
||
|
|
if required not in nodes:
|
||
|
|
raise ValueError(f"required STRING tree node missing: {required}")
|
||
|
|
|
||
|
|
tree_string_taxa = {
|
||
|
|
taxon_id for taxon_id, node in nodes.items() if node.is_string_species
|
||
|
|
}
|
||
|
|
missing_from_tree = sorted(registry_taxa - tree_string_taxa, key=int)
|
||
|
|
missing_from_registry = sorted(tree_string_taxa - registry_taxa, key=int)
|
||
|
|
if missing_from_tree or missing_from_registry:
|
||
|
|
raise ValueError(
|
||
|
|
"STRING registry/tree STRING species mismatch: "
|
||
|
|
f"registry_only={missing_from_tree[:10]}, tree_only={missing_from_registry[:10]}"
|
||
|
|
)
|
||
|
|
if "3702" not in registry_taxa:
|
||
|
|
raise ValueError("Arabidopsis 3702 is not a STRING species")
|
||
|
|
|
||
|
|
return {
|
||
|
|
"viridiplantae_taxon_id": "33090",
|
||
|
|
"arabidopsis_taxon_id": "3702",
|
||
|
|
"registry_species_count": len(registry_taxa),
|
||
|
|
"tree_node_count": len(nodes),
|
||
|
|
"string_species_count": len(tree_string_taxa),
|
||
|
|
"root_taxon_id": "1",
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _download(url: str, destination: Path) -> None:
|
||
|
|
request = urllib.request.Request(url, headers={"User-Agent": "PlantAI-STRING-snapshot/1"})
|
||
|
|
with urllib.request.urlopen(request, timeout=300) as response:
|
||
|
|
with destination.open("wb") as output:
|
||
|
|
shutil.copyfileobj(response, output)
|
||
|
|
|
||
|
|
|
||
|
|
def prepare_snapshot(
|
||
|
|
destination: Path,
|
||
|
|
version: str = STRING_VERSION,
|
||
|
|
source_directory: Path | None = None,
|
||
|
|
) -> dict[str, object]:
|
||
|
|
if version != STRING_VERSION:
|
||
|
|
raise ValueError(f"only STRING version {STRING_VERSION} is supported")
|
||
|
|
destination = Path(destination)
|
||
|
|
destination.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
with tempfile.TemporaryDirectory(prefix="string-snapshot-", dir=destination) as tmp:
|
||
|
|
temporary = Path(tmp)
|
||
|
|
for name, url in SOURCE_URLS.items():
|
||
|
|
if source_directory is None:
|
||
|
|
_download(url, temporary / name)
|
||
|
|
else:
|
||
|
|
source = Path(source_directory) / name
|
||
|
|
if not source.is_file():
|
||
|
|
raise ValueError(f"local STRING source file missing: {source}")
|
||
|
|
shutil.copy2(source, temporary / name)
|
||
|
|
validation = validate_snapshot(
|
||
|
|
temporary / "species.v12.0.txt",
|
||
|
|
temporary / "species.tree.v12.0.txt",
|
||
|
|
)
|
||
|
|
for name in SOURCE_URLS:
|
||
|
|
os.replace(temporary / name, destination / name)
|
||
|
|
|
||
|
|
files = {
|
||
|
|
name: {
|
||
|
|
"source_url": url,
|
||
|
|
"size": (destination / name).stat().st_size,
|
||
|
|
"sha256": _sha256(destination / name),
|
||
|
|
}
|
||
|
|
for name, url in SOURCE_URLS.items()
|
||
|
|
}
|
||
|
|
snapshot_id = hashlib.sha256(
|
||
|
|
(
|
||
|
|
PARSER_SCHEMA_VERSION
|
||
|
|
+ "\n"
|
||
|
|
+ files["species.v12.0.txt"]["sha256"]
|
||
|
|
+ "\n"
|
||
|
|
+ files["species.tree.v12.0.txt"]["sha256"]
|
||
|
|
).encode("utf-8")
|
||
|
|
).hexdigest()
|
||
|
|
manifest: dict[str, object] = {
|
||
|
|
"string_version": version,
|
||
|
|
"source_urls": dict(SOURCE_URLS),
|
||
|
|
"transfer_mode": "local_files" if source_directory is not None else "download",
|
||
|
|
"downloaded_at": datetime.now(timezone.utc).isoformat(),
|
||
|
|
"parser_schema_version": PARSER_SCHEMA_VERSION,
|
||
|
|
"snapshot_id": snapshot_id,
|
||
|
|
"files": files,
|
||
|
|
"validation": validation,
|
||
|
|
}
|
||
|
|
manifest_path = destination / "source_manifest.json"
|
||
|
|
with tempfile.NamedTemporaryFile(
|
||
|
|
mode="w", encoding="utf-8", dir=destination, delete=False
|
||
|
|
) as handle:
|
||
|
|
json.dump(manifest, handle, ensure_ascii=False, indent=2)
|
||
|
|
handle.write("\n")
|
||
|
|
temporary_manifest = Path(handle.name)
|
||
|
|
os.replace(temporary_manifest, manifest_path)
|
||
|
|
return manifest
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> int:
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--version", default=STRING_VERSION)
|
||
|
|
parser.add_argument("--destination", type=Path, required=True)
|
||
|
|
parser.add_argument("--source-directory", type=Path)
|
||
|
|
args = parser.parse_args()
|
||
|
|
manifest = prepare_snapshot(
|
||
|
|
args.destination,
|
||
|
|
version=args.version,
|
||
|
|
source_directory=args.source_directory,
|
||
|
|
)
|
||
|
|
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
raise SystemExit(main())
|