diff --git a/.gitignore b/.gitignore index 3139ede..a3e9114 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,8 @@ __pycache__ -_build \ No newline at end of file +_build +build/ +*.egg-info/ + +# Legacy in-package taxonomy cache written by bkbit < this release; the cache +# now lives in a per-user cache dir (see bkbit/utils/ncbi_taxonomy_cache.py). +bkbit/utils/ncbi_taxonomy/ diff --git a/CLAUDE.md b/CLAUDE.md index c8079c3..dbf5339 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,7 +47,7 @@ The package provides a CLI tool `bkbit` with multiple subcommands: - `bkbit schema2model` - Convert spreadsheets to LinkML YAML models - `bkbit yaml2cvs` - Convert YAML to CSV - `bkbit linkml-trimmer` - Trim LinkML models -- `bkbit download-ncbi-taxonomy` - Download NCBI taxonomy data +- `bkbit download-ncbi-taxonomy` - Optionally pre-download the full NCBI taxonomy (see below) ## Architecture @@ -73,7 +73,9 @@ The package provides a CLI tool `bkbit` with multiple subcommands: - `add_dunderMethods_genomeAnnotation.py` - Adds __eq__, __ne__, __hash__ to GeneAnnotation **`bkbit/utils/`** - Shared utilities: -- `get_ncbi_taxonomy.py` - NCBI taxonomy download +- `ncbi_taxonomy_cache.py` - NCBI taxonomy name lookups (see NCBI Taxonomy Data below) +- `ncbi_taxonomy_data/` - Bundled taxonomy subset + the script that rebuilds it +- `get_ncbi_taxonomy.py` - Thin CLI wrapper for pre-fetching the full taxonomy - `nimp_api_endpoints.py` - Specimen Portal API endpoints - `setup_logger.py` - Logging configuration @@ -82,5 +84,31 @@ The package provides a CLI tool `bkbit` with multiple subcommands: - **Schemasheets** - Spreadsheet to LinkML conversion - **Click** - CLI framework +### NCBI Taxonomy Data + +`gff2jsonld` resolves organism names through `bkbit/utils/ncbi_taxonomy_cache.py`, which +serves lookups from two layers and requires no setup after `pip install`: + +1. **Bundled subset** (`bkbit/utils/ncbi_taxonomy_data/taxonomy_subset.json.gz`, ~575KB) - + every taxon with a GenBank common name (~30k), which covers all BICAN organisms. + Offline, loads in milliseconds. +2. **Full NCBI taxonomy** - downloaded on demand only if a lookup misses the subset, and + cached in a per-user cache directory (never in `site-packages`). + +Rules to preserve when touching this code: +- **Never load taxonomy data at import time.** The previous version read the maps in the + `Gff3` class body, so a missing cache broke every `bkbit` subcommand, including `--help`. +- **Never write runtime data into the installed package.** `site-packages` may be read-only + and is wiped on upgrade. + +To regenerate the bundled subset for a newer taxonomy dump: +```bash +bkbit download-ncbi-taxonomy +python -m bkbit.utils.ncbi_taxonomy_data.build_subset # commit the resulting .json.gz +``` + ### Environment Variables - `jwt_token` - Specimen Portal Personal API Token (required for specimen2jsonld) +- `BKBIT_DATA_DIR` - Overrides the cache directory for the downloaded full NCBI taxonomy +- `BKBIT_NO_DOWNLOAD` - If set, a taxonomy lookup that misses the bundled subset raises + instead of downloading (useful in CI/air-gapped runs) diff --git a/bkbit/data_translators/HMBA_annotation_translator.py b/bkbit/data_translators/HMBA_annotation_translator.py index f061760..2929a6a 100644 --- a/bkbit/data_translators/HMBA_annotation_translator.py +++ b/bkbit/data_translators/HMBA_annotation_translator.py @@ -1,6 +1,5 @@ import itertools import click -import pkg_resources from bkbit.models import bke_taxonomy import json import hashlib diff --git a/bkbit/data_translators/genome_annotation_translator.py b/bkbit/data_translators/genome_annotation_translator.py index 54f7a0c..72512bc 100644 --- a/bkbit/data_translators/genome_annotation_translator.py +++ b/bkbit/data_translators/genome_annotation_translator.py @@ -41,12 +41,11 @@ import gzip from tqdm import tqdm import click -import pkg_resources from linkml_runtime.dumpers import json_dumper from rdflib import Graph from bkbit.models import genome_annotation as ga from bkbit.utils.setup_logger import setup_logger -from bkbit.utils.load_json import load_json +from bkbit.utils import ncbi_taxonomy_cache from bkbit.utils.generate_bkbit_id import generate_object_id from bkbit.utils.serialize_to_ttl import convert_jsonld_to_ttl @@ -72,10 +71,6 @@ ) DEFAULT_FEATURE_FILTER = ("gene", "pseudogene", "ncRNA_gene") DEFAULT_HASH = ("MD5",) -TAXON_DIR_PATH = "../utils/ncbi_taxonomy/" -SCIENTIFIC_NAME_TO_TAXONID_PATH = pkg_resources.resource_filename(__name__, TAXON_DIR_PATH + "scientific_name_to_taxid.json") -TAXON_SCIENTIFIC_NAME_PATH = pkg_resources.resource_filename(__name__, TAXON_DIR_PATH + "taxid_to_scientific_name.json") -TAXON_COMMON_NAME_PATH = pkg_resources.resource_filename(__name__, TAXON_DIR_PATH + "taxid_to_common_name.json") INCOMPATABLE_EXTENSION = "The provided content URL is not supported. Please provide a valid URL with '.gff.gz' extension." class Gff3: """ @@ -134,18 +129,10 @@ class Gff3: serialize_to_jsonld(exclude_none=True, exclude_unset=False): Serializes the object and either writes it to the specified output file or prints it to the CLI. """ - scientific_name_to_taxonid = None - taxon_scientific_name = None - taxon_common_name = None - - # Load taxonomy data at class definition - try: - scientific_name_to_taxonid = load_json(SCIENTIFIC_NAME_TO_TAXONID_PATH) - taxon_scientific_name = load_json(TAXON_SCIENTIFIC_NAME_PATH) - taxon_common_name = load_json(TAXON_COMMON_NAME_PATH) - except FileNotFoundError as e: - #logging.critical("NCBI Taxonomy not downloaded. Run 'bkbit download-ncbi-taxonomy' first.") - raise RuntimeError("NCBI Taxonomy not downloaded. Run 'bkbit download-ncbi-taxonomy' first.") from e + # Taxonomy names are resolved lazily, one taxon at a time, by + # bkbit.utils.ncbi_taxonomy_cache. Nothing is loaded or downloaded until a + # lookup actually happens, so importing this module is always side-effect + # free. def __init__( self, @@ -299,7 +286,7 @@ def parse_url(self, assembly_accession: str = None): ) scientific_name = ensembl_match.group(3) - taxonid = self.scientific_name_to_taxonid.get( + taxonid = ncbi_taxonomy_cache.lookup_taxid( scientific_name.replace("_", " ") ) return { @@ -377,8 +364,18 @@ def generate_organism_taxon(cls, taxon_id: str): Returns: ga.OrganismTaxon: The generated organism taxon object. + + Raises: + ValueError: If the taxon ID is not present in the NCBI taxonomy. """ - attributes = {"full_name": cls.taxon_scientific_name[taxon_id], "name": cls.taxon_common_name[taxon_id], "iri": PREFIX_MAP[TAXON_PREFIX] + taxon_id, "xref": [TAXON_PREFIX + taxon_id]} + scientific_name = ncbi_taxonomy_cache.lookup_scientific_name(taxon_id) + if scientific_name is None: + raise ValueError( + f"Taxon ID '{taxon_id}' was not found in the NCBI taxonomy." + ) + # Not every taxon has a GenBank common name; fall back to the scientific name. + common_name = ncbi_taxonomy_cache.lookup_common_name(taxon_id) or scientific_name + attributes = {"full_name": scientific_name, "name": common_name, "iri": PREFIX_MAP[TAXON_PREFIX] + taxon_id, "xref": [TAXON_PREFIX + taxon_id]} attributes["id"] = generate_object_id(attributes) return ga.OrganismTaxon(**attributes) diff --git a/bkbit/model_editors/add_dunderMethods_genomeAnnotation.py b/bkbit/model_editors/add_dunderMethods_genomeAnnotation.py index 27e9501..16287e9 100644 --- a/bkbit/model_editors/add_dunderMethods_genomeAnnotation.py +++ b/bkbit/model_editors/add_dunderMethods_genomeAnnotation.py @@ -1,7 +1,7 @@ import re -import pkg_resources +from pathlib import Path # Read the file -genome_annotation_model = pkg_resources.resource_filename(__name__, "../models/genome_annotation.py") +genome_annotation_model = Path(__file__).parent.parent / "models" / "genome_annotation.py" #file_path = "bkbit/models/genome_annotation.py" with open(genome_annotation_model, "r") as file: diff --git a/bkbit/utils/get_ncbi_taxonomy.py b/bkbit/utils/get_ncbi_taxonomy.py index 2b57023..2ab5610 100644 --- a/bkbit/utils/get_ncbi_taxonomy.py +++ b/bkbit/utils/get_ncbi_taxonomy.py @@ -1,187 +1,72 @@ """ -This script downloads a zip file containing taxonomic data from a given URL, extracts and processes -the content of the 'names.dmp' file in memory, and saves the parsed data into JSON files. The script -includes three main functions: +CLI for pre-fetching the full NCBI taxonomy cache. -1. download_and_extract_zip_in_memory(url): - Downloads a zip file from the given URL and extracts the content of the 'names.dmp' file in memory. +bkbit does not require this step: a taxonomy subset covering every organism with +a GenBank common name ships inside the wheel, and the full taxonomy is downloaded +automatically the first time a lookup falls outside that subset. This command +exists for the cases where you want the download to happen up front rather than +mid-pipeline - air-gapped runs, container images, CI, or reproducible builds. -2. parse_dmp_content(dmp_content): - Parses the content of a DMP file and extracts taxonomic information into dictionaries. - -3. process_and_save_taxdmp_in_memory(url, output_dir): - Downloads and processes the taxdump file from the given URL, and saves the parsed data into - separate JSON files in the specified output directory. +All of the download and cache logic lives in +:mod:`bkbit.utils.ncbi_taxonomy_cache`; this module is a thin wrapper around it +and re-exports the previous function names for backwards compatibility. Usage: - The script can be executed as a standalone program. Modify the URL and output directory as needed. + bkbit download-ncbi-taxonomy [--reload] [--data-dir PATH] """ -import json -import zipfile -import io -import os -import requests -import pkg_resources import click -NCBI_TAXON_URL = "https://ftp.ncbi.nih.gov/pub/taxonomy/taxdmp.zip" -OUTPUT_DIR_NAME = "ncbi_taxonomy" -OUTPUT_DIR_PATH = pkg_resources.resource_filename(__name__, OUTPUT_DIR_NAME) -SCIENTIFIC_NAME_TO_TAXONID_PATH = pkg_resources.resource_filename(__name__, "ncbi_taxonomy/scientific_name_to_taxid.json") -TAXON_SCIENTIFIC_NAME_PATH = pkg_resources.resource_filename(__name__, "ncbi_taxonomy/taxid_to_scientific_name.json") -TAXON_COMMON_NAME_PATH = pkg_resources.resource_filename(__name__, "ncbi_taxonomy/taxid_to_common_name.json") - - - -def download_and_extract_zip_in_memory(url): +from bkbit.utils.ncbi_taxonomy_cache import ( + NCBI_TAXON_URL, + build_full_cache, + download_and_extract_zip_in_memory, + ensure_full_cache, + parse_dmp_content, +) + +__all__ = [ + "NCBI_TAXON_URL", + "build_full_cache", + "download_and_extract_zip_in_memory", + "download_ncbi_taxonomy", + "parse_dmp_content", + "process_and_save_taxdmp_in_memory", +] + + +def process_and_save_taxdmp_in_memory(url=NCBI_TAXON_URL, output_dir=None): """ - Downloads a zip file from the given URL and extracts the content of the 'names.dmp' file in memory. - - Args: - url (str): The URL of the zip file to download. - - Returns: - str: The content of the 'names.dmp' file as a string. - - Raises: - requests.exceptions.HTTPError: If the file download fails with a non-200 status code. - """ - # Download the file - response = requests.get(url, timeout=30) - if response.status_code == 200: - # Unzip the file in memory - with zipfile.ZipFile(io.BytesIO(response.content)) as z: - # Extract names.dmp file content into memory - with z.open("names.dmp") as names_dmp_file: - names_dmp_content = names_dmp_file.read().decode("utf-8") - return names_dmp_content - else: - raise requests.exceptions.HTTPError( - f"Failed to download file, status code: {response.status_code}" - ) - + Downloads and processes the taxdump file, saving the parsed data as JSON. -def parse_dmp_content(dmp_content): - """ - Parses the content of a DMP file and extracts taxonomic information. + Deprecated alias for :func:`bkbit.utils.ncbi_taxonomy_cache.build_full_cache`. Args: - dmp_content (str): The content of the DMP file. - - Returns: - tuple: A tuple containing three dictionaries: - - taxid_to_scientific_name: A dictionary mapping taxonomic IDs to scientific names. - - taxid_to_common_name: A dictionary mapping taxonomic IDs to common names. - - scientific_name_to_taxid: A dictionary mapping scientific names to taxonomic IDs. + url: The URL of the taxdump file to download and process. + output_dir: Destination directory. Defaults to the bkbit cache directory. """ - taxid_to_scientific_name = {} - taxid_to_common_name = {} - scientific_name_to_taxid = {} - - for line in dmp_content.strip().split("\n"): - # Split the line by the delimiter '|' - parts = line.strip().split("|") - - # Remove leading and trailing whitespace from each part - parts = [part.strip() for part in parts] - # Taxonomy names file (names.dmp): - # tax_id-- the id of node associated with this name - # name_txt-- name itself - # unique name-- the unique variant of this name if name not unique - # name class-- (synonym, common name, ...) - taxid = parts[0] - name = parts[1] - unique_name = parts[2] - name_class = parts[3] - - # Create a dictionary with the parsed data - if name_class == "scientific name" and taxid not in taxid_to_scientific_name: - if unique_name: - taxid_to_scientific_name[taxid] = unique_name - scientific_name_to_taxid[unique_name] = taxid - else: - taxid_to_scientific_name[taxid] = name - scientific_name_to_taxid[name] = taxid - elif name_class == "genbank common name" and taxid not in taxid_to_common_name: - taxid_to_common_name[taxid] = name - return taxid_to_scientific_name, taxid_to_common_name, scientific_name_to_taxid - + build_full_cache(url=url, cache_dir=output_dir) -def process_and_save_taxdmp_in_memory(url, output_dir): - """ - Downloads and processes the taxdump file from the given URL, - and saves the parsed data into separate JSON files in the specified output directory. - - Args: - url (str): The URL of the taxdump file to download and process. - output_dir (str): The directory where the parsed data will be saved. - - Returns: - None - """ - # Ensure the output directory exists - if not os.path.exists(output_dir): - os.makedirs(output_dir) - - # Step 1: Download and unzip the folder in memory - names_dmp_content = download_and_extract_zip_in_memory(url) - - # Step 2: Parse the names.dmp content - taxid_to_scientific_name, taxid_to_common_name, scientific_name_to_taxid = ( - parse_dmp_content(names_dmp_content) - ) - - # Step 3: Save the dictionaries to files - with open( - os.path.join(output_dir, "taxid_to_common_name.json"), "w", encoding="utf-8" - ) as f: - json.dump(taxid_to_common_name, f, indent=4) - - with open( - os.path.join(output_dir, "taxid_to_scientific_name.json"), "w", encoding="utf-8" - ) as f: - json.dump(taxid_to_scientific_name, f, indent=4) - - with open( - os.path.join(output_dir, "scientific_name_to_taxid.json"), "w", encoding="utf-8" - ) as f: - json.dump(scientific_name_to_taxid, f, indent=4) - - - -def load_json(file_path): - """ - Load JSON data from a file. - - Args: - file_path (str): The path to the JSON file. - - Returns: - dict: The loaded JSON data. - - """ - with open(file_path, "r", encoding="utf-8") as f: - return json.load(f) @click.command() -@click.option("--reload", '-r', is_flag=True, help="Reload NCBI taxonomy data") - -def download_ncbi_taxonomy(reload=False): - - """ - Load JSON data from a file. - - Args: - file_path (str): The path to the JSON file. - - Returns: - dict: The loaded JSON data as a dictionary. - """ - if reload or not os.path.exists(SCIENTIFIC_NAME_TO_TAXONID_PATH) or not os.path.exists(TAXON_SCIENTIFIC_NAME_PATH) or not os.path.exists(TAXON_COMMON_NAME_PATH): - process_and_save_taxdmp_in_memory(NCBI_TAXON_URL, OUTPUT_DIR_PATH) +@click.option( + "--reload", "-r", is_flag=True, help="Re-download even if already cached." +) +@click.option( + "--data-dir", + "-d", + type=click.Path(file_okay=False), + default=None, + help="Directory to store the taxonomy cache in. Defaults to the bkbit cache directory (override with BKBIT_DATA_DIR).", +) +def download_ncbi_taxonomy(reload=False, data_dir=None): + """Pre-download the full NCBI taxonomy used by gff2jsonld (optional).""" + target = ensure_full_cache(reload=reload, cache_dir=data_dir) + if reload: + click.echo(f"Re-downloaded NCBI taxonomy to {target}") else: - print("PRINT already downloaded") + click.echo(f"NCBI taxonomy cache ready at {target}") + if __name__ == "__main__": - download_ncbi_taxonomy() + download_ncbi_taxonomy() diff --git a/bkbit/utils/ncbi_taxonomy_cache.py b/bkbit/utils/ncbi_taxonomy_cache.py new file mode 100644 index 0000000..585e2d3 --- /dev/null +++ b/bkbit/utils/ncbi_taxonomy_cache.py @@ -0,0 +1,312 @@ +""" +Lazy, self-provisioning access to NCBI taxonomy names. + +bkbit needs three lookups from the NCBI taxonomy: scientific name -> taxon id, +taxon id -> scientific name, and taxon id -> common name. This module serves +those lookups from two layers, in order: + +1. A subset bundled inside the wheel (`ncbi_taxonomy_data/taxonomy_subset.json.gz`, + ~600KB, every taxon that has a GenBank common name). No network, no setup, + loads in milliseconds. This covers effectively every organism a genome + annotation pipeline is run against. +2. The full NCBI taxonomy dump, downloaded and cached on first use for taxa that + are not in the subset. The cache lives in a per-user cache directory - never + inside `site-packages`, which may be read-only and is wiped on upgrade. + +Nothing here runs at import time: the bundled subset is loaded on first lookup +and the full dump is only fetched if a lookup actually misses the subset. + +Environment variables: + BKBIT_DATA_DIR: Overrides the cache directory used for the full taxonomy. + BKBIT_NO_DOWNLOAD: If set to a truthy value, a lookup that misses the + bundled subset raises instead of downloading the full taxonomy. +""" + +import gzip +import io +import json +import os +import zipfile +from functools import lru_cache +from importlib.resources import files +from pathlib import Path +from typing import Dict, Optional, Tuple + +import platformdirs +import requests + +NCBI_TAXON_URL = "https://ftp.ncbi.nih.gov/pub/taxonomy/taxdmp.zip" + +BUNDLED_SUBSET_PACKAGE = "bkbit.utils.ncbi_taxonomy_data" +BUNDLED_SUBSET_FILENAME = "taxonomy_subset.json.gz" +SUBSET_FORMAT_VERSION = 1 + +DATA_DIR_ENV_VAR = "BKBIT_DATA_DIR" +NO_DOWNLOAD_ENV_VAR = "BKBIT_NO_DOWNLOAD" + +CACHE_SUBDIR = "ncbi_taxonomy" +CACHE_FILENAMES = { + "scientific": "taxid_to_scientific_name.json", + "common": "taxid_to_common_name.json", + "scientific_to_taxid": "scientific_name_to_taxid.json", +} + +# Older bkbit versions wrote the cache into the installed package directory. +# Reused read-only if it happens to still be there, but never written to. +LEGACY_CACHE_DIR = Path(__file__).parent / CACHE_SUBDIR + + +## CACHE LOCATION ## + + +def data_dir() -> Path: + """ + Returns the directory bkbit uses for the downloaded NCBI taxonomy cache. + + Honours ``BKBIT_DATA_DIR`` if set, otherwise falls back to the platform's + per-user cache directory (which respects ``XDG_CACHE_HOME`` on Linux). + """ + override = os.environ.get(DATA_DIR_ENV_VAR) + if override: + return Path(override).expanduser() + return Path(platformdirs.user_cache_dir("bkbit", appauthor=False)) / CACHE_SUBDIR + + +def full_cache_paths(cache_dir: Optional[Path] = None) -> Dict[str, Path]: + """ + Returns the paths of the three JSON files that make up the full cache. + + Args: + cache_dir: Directory to resolve against. Defaults to :func:`data_dir`. + + Returns: + dict: Keys ``scientific``, ``common``, and ``scientific_to_taxid``. + """ + base = Path(cache_dir) if cache_dir is not None else data_dir() + return {key: base / name for key, name in CACHE_FILENAMES.items()} + + +def _existing_cache_dir() -> Optional[Path]: + """ + Returns a directory containing a complete taxonomy cache, or None. + + Prefers the current cache location and falls back to the legacy in-package + directory written by older bkbit versions. + """ + candidates = [data_dir()] + if not os.environ.get(DATA_DIR_ENV_VAR): + candidates.append(LEGACY_CACHE_DIR) + for candidate in candidates: + if all(path.exists() for path in full_cache_paths(candidate).values()): + return candidate + return None + + +## DOWNLOAD / BUILD ## + + +def download_and_extract_zip_in_memory(url: str = NCBI_TAXON_URL) -> str: + """ + Downloads the taxdump zip from the given URL and returns 'names.dmp'. + + Args: + url: The URL of the zip file to download. + + Returns: + str: The content of the 'names.dmp' file as a string. + + Raises: + requests.exceptions.HTTPError: If the download fails. + """ + response = requests.get(url, timeout=300) + response.raise_for_status() + with zipfile.ZipFile(io.BytesIO(response.content)) as z: + with z.open("names.dmp") as names_dmp_file: + return names_dmp_file.read().decode("utf-8") + + +def parse_dmp_content(dmp_content: str) -> Tuple[dict, dict, dict]: + """ + Parses the content of a names.dmp file into taxonomy name lookups. + + Args: + dmp_content: The content of the DMP file. + + Returns: + tuple: ``(taxid_to_scientific_name, taxid_to_common_name, + scientific_name_to_taxid)``. + """ + taxid_to_scientific_name = {} + taxid_to_common_name = {} + scientific_name_to_taxid = {} + + for line in dmp_content.strip().split("\n"): + parts = [part.strip() for part in line.strip().split("|")] + # names.dmp columns: tax_id | name_txt | unique name | name class + taxid, name, unique_name, name_class = parts[0], parts[1], parts[2], parts[3] + + if name_class == "scientific name" and taxid not in taxid_to_scientific_name: + resolved = unique_name or name + taxid_to_scientific_name[taxid] = resolved + scientific_name_to_taxid[resolved] = taxid + elif name_class == "genbank common name" and taxid not in taxid_to_common_name: + taxid_to_common_name[taxid] = name + + return taxid_to_scientific_name, taxid_to_common_name, scientific_name_to_taxid + + +def build_full_cache( + url: str = NCBI_TAXON_URL, cache_dir: Optional[Path] = None +) -> Dict[str, Path]: + """ + Downloads the NCBI taxdump and writes the full cache to disk. + + Args: + url: The URL of the taxdump zip to download and process. + cache_dir: Destination directory. Defaults to :func:`data_dir`. + + Returns: + dict: The written cache paths, as returned by :func:`full_cache_paths`. + """ + paths = full_cache_paths(cache_dir) + target_dir = next(iter(paths.values())).parent + target_dir.mkdir(parents=True, exist_ok=True) + + names_dmp_content = download_and_extract_zip_in_memory(url) + scientific, common, scientific_to_taxid = parse_dmp_content(names_dmp_content) + + # Written compactly; these maps run to hundreds of MB when pretty-printed. + for key, payload in ( + ("scientific", scientific), + ("common", common), + ("scientific_to_taxid", scientific_to_taxid), + ): + with paths[key].open("w", encoding="utf-8") as f: + json.dump(payload, f, separators=(",", ":"), ensure_ascii=False) + + return paths + + +def ensure_full_cache(reload: bool = False, cache_dir: Optional[Path] = None) -> Path: + """ + Makes sure the full taxonomy cache exists on disk, downloading if needed. + + Args: + reload: Re-download even if a complete cache is already present. + cache_dir: Destination directory. Defaults to :func:`data_dir`. + + Returns: + Path: The directory containing the cache. + """ + if not reload: + if cache_dir is None: + existing = _existing_cache_dir() + if existing is not None: + return existing + elif all(path.exists() for path in full_cache_paths(cache_dir).values()): + return Path(cache_dir) + + paths = build_full_cache(cache_dir=cache_dir) + _full_taxonomy.cache_clear() + return next(iter(paths.values())).parent + + +## LOOKUPS ## + + +@lru_cache(maxsize=1) +def _bundled_taxa() -> Dict[str, list]: + """ + Loads the bundled taxonomy subset: ``{taxid: [scientific, common]}``. + """ + raw = (files(BUNDLED_SUBSET_PACKAGE) / BUNDLED_SUBSET_FILENAME).read_bytes() + payload = json.loads(gzip.decompress(raw).decode("utf-8")) + if payload.get("format") != SUBSET_FORMAT_VERSION: + raise RuntimeError( + f"Unsupported bundled taxonomy subset format: {payload.get('format')!r}" + ) + return payload["taxa"] + + +@lru_cache(maxsize=1) +def _bundled_scientific_name_to_taxid() -> Dict[str, str]: + """ + Inverts the bundled subset into a scientific name -> taxon id lookup. + """ + return {names[0]: taxid for taxid, names in _bundled_taxa().items()} + + +@lru_cache(maxsize=1) +def _full_taxonomy() -> Tuple[dict, dict, dict]: + """ + Loads the full taxonomy from the cache, downloading it if necessary. + + Returns: + tuple: ``(taxid_to_scientific_name, taxid_to_common_name, + scientific_name_to_taxid)``. + + Raises: + RuntimeError: If the cache is missing and downloads are disabled via + ``BKBIT_NO_DOWNLOAD``. + """ + cache_dir = _existing_cache_dir() + if cache_dir is None: + if os.environ.get(NO_DOWNLOAD_ENV_VAR): + raise RuntimeError( + "This taxon is not in the taxonomy subset bundled with bkbit and " + f"{NO_DOWNLOAD_ENV_VAR} is set. Run 'bkbit download-ncbi-taxonomy' " + f"or unset {NO_DOWNLOAD_ENV_VAR} to allow the download." + ) + print( + "Taxon not found in the bundled NCBI taxonomy subset; downloading the " + f"full NCBI taxonomy to {data_dir()} (this happens once)." + ) + cache_dir = ensure_full_cache() + + paths = full_cache_paths(cache_dir) + with paths["scientific"].open(encoding="utf-8") as f: + scientific = json.load(f) + with paths["common"].open(encoding="utf-8") as f: + common = json.load(f) + with paths["scientific_to_taxid"].open(encoding="utf-8") as f: + scientific_to_taxid = json.load(f) + return scientific, common, scientific_to_taxid + + +def lookup_taxid(scientific_name: str) -> Optional[str]: + """ + Returns the taxon id for a scientific name, or None if it is unknown. + + Args: + scientific_name: Scientific name, e.g. ``"Homo sapiens"``. + """ + taxid = _bundled_scientific_name_to_taxid().get(scientific_name) + if taxid is not None: + return taxid + return _full_taxonomy()[2].get(scientific_name) + + +def lookup_scientific_name(taxid: str) -> Optional[str]: + """ + Returns the scientific name for a taxon id, or None if it is unknown. + + Args: + taxid: NCBI taxon id, as a string, e.g. ``"9606"``. + """ + names = _bundled_taxa().get(str(taxid)) + if names is not None: + return names[0] + return _full_taxonomy()[0].get(str(taxid)) + + +def lookup_common_name(taxid: str) -> Optional[str]: + """ + Returns the GenBank common name for a taxon id, or None if it has none. + + Args: + taxid: NCBI taxon id, as a string, e.g. ``"9606"``. + """ + names = _bundled_taxa().get(str(taxid)) + if names is not None: + return names[1] + return _full_taxonomy()[1].get(str(taxid)) diff --git a/bkbit/utils/ncbi_taxonomy_data/__init__.py b/bkbit/utils/ncbi_taxonomy_data/__init__.py new file mode 100644 index 0000000..ceb5764 --- /dev/null +++ b/bkbit/utils/ncbi_taxonomy_data/__init__.py @@ -0,0 +1 @@ +"""Bundled NCBI taxonomy subset shipped with bkbit (see build_subset.py).""" diff --git a/bkbit/utils/ncbi_taxonomy_data/build_subset.py b/bkbit/utils/ncbi_taxonomy_data/build_subset.py new file mode 100644 index 0000000..0664b76 --- /dev/null +++ b/bkbit/utils/ncbi_taxonomy_data/build_subset.py @@ -0,0 +1,93 @@ +""" +Rebuilds the NCBI taxonomy subset that is bundled inside the bkbit wheel. + +The full NCBI taxonomy (~2.6M taxa, ~230MB of JSON) is far too large to ship on +PyPI, but bkbit only ever needs a scientific name, a common name, and a taxon id +for the organism a GFF3 file belongs to. Every taxon that carries a GenBank +common name (~30k of them, which covers every organism anyone realistically +runs an annotation pipeline for) fits in well under 1MB gzipped, so that subset +is bundled and used as the fast, offline path. Anything outside the subset falls +back to the full download managed by `bkbit.utils.ncbi_taxonomy_cache`. + +Usage (maintainers only, not part of the runtime path): + + bkbit download-ncbi-taxonomy # build the full cache first + python -m bkbit.utils.ncbi_taxonomy_data.build_subset + +Commit the regenerated `taxonomy_subset.json.gz` alongside any release that +should pick up a newer taxonomy dump. +""" + +import gzip +import json +from datetime import date +from pathlib import Path +from typing import Optional + +from bkbit.utils.ncbi_taxonomy_cache import ( + BUNDLED_SUBSET_FILENAME, + NCBI_TAXON_URL, + SUBSET_FORMAT_VERSION, + full_cache_paths, +) + + +def build_subset( + cache_dir: Optional[Path] = None, output_path: Optional[Path] = None +) -> Path: + """ + Builds the bundled subset from a fully downloaded taxonomy cache. + + Args: + cache_dir: Directory holding the full taxonomy cache. Defaults to the + cache directory `bkbit.utils.ncbi_taxonomy_cache` resolves to. + output_path: Where to write the gzipped subset. Defaults to the bundled + location inside this package. + + Returns: + Path: The path the subset was written to. + """ + paths = full_cache_paths(cache_dir) + missing = [str(p) for p in paths.values() if not p.exists()] + if missing: + raise FileNotFoundError( + "Full taxonomy cache is incomplete; run 'bkbit download-ncbi-taxonomy' " + f"first. Missing: {', '.join(missing)}" + ) + + with paths["scientific"].open(encoding="utf-8") as f: + taxid_to_scientific_name = json.load(f) + with paths["common"].open(encoding="utf-8") as f: + taxid_to_common_name = json.load(f) + + # A taxon is only usable by the translator if it has both names, so the + # subset is exactly the intersection of the two maps. + taxa = { + taxid: [taxid_to_scientific_name[taxid], common_name] + for taxid, common_name in taxid_to_common_name.items() + if taxid in taxid_to_scientific_name + } + + payload = { + "format": SUBSET_FORMAT_VERSION, + "source": NCBI_TAXON_URL, + "built": date.today().isoformat(), + "taxa": taxa, + } + + if output_path is None: + output_path = Path(__file__).parent / BUNDLED_SUBSET_FILENAME + blob = json.dumps(payload, separators=(",", ":"), ensure_ascii=False).encode( + "utf-8" + ) + with gzip.open(output_path, "wb", compresslevel=9) as f: + f.write(blob) + + print( + f"Wrote {len(taxa)} taxa to {output_path} ({output_path.stat().st_size} bytes)" + ) + return output_path + + +if __name__ == "__main__": + build_subset() diff --git a/bkbit/utils/ncbi_taxonomy_data/taxonomy_subset.json.gz b/bkbit/utils/ncbi_taxonomy_data/taxonomy_subset.json.gz new file mode 100644 index 0000000..a6ed804 Binary files /dev/null and b/bkbit/utils/ncbi_taxonomy_data/taxonomy_subset.json.gz differ diff --git a/docs/bkbit.utils.ncbi_taxonomy_cache.rst b/docs/bkbit.utils.ncbi_taxonomy_cache.rst new file mode 100644 index 0000000..8a7c303 --- /dev/null +++ b/docs/bkbit.utils.ncbi_taxonomy_cache.rst @@ -0,0 +1,7 @@ +bkbit.utils.ncbi\_taxonomy\_cache module +======================================== + +.. automodule:: bkbit.utils.ncbi_taxonomy_cache + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/bkbit.utils.rst b/docs/bkbit.utils.rst index a609e9a..65af97f 100644 --- a/docs/bkbit.utils.rst +++ b/docs/bkbit.utils.rst @@ -9,6 +9,7 @@ Submodules bkbit.utils.get_ncbi_taxonomy bkbit.utils.load_json + bkbit.utils.ncbi_taxonomy_cache bkbit.utils.nimp_api_endpoints bkbit.utils.setup_logger diff --git a/docs/genome_annotation.rst b/docs/genome_annotation.rst index 9d941be..20e6c9d 100644 --- a/docs/genome_annotation.rst +++ b/docs/genome_annotation.rst @@ -16,7 +16,26 @@ Each JSON-LD file will contain: - 1 OrganismTaxon object - 1 Checksum object -Command Line +NCBI Taxonomy Data +................... + +No setup is required before running ``bkbit gff2jsonld``. A taxonomy subset covering every +organism with a GenBank common name ships inside the ``bkbit`` package, so organism names +resolve offline for all supported species. + +If a GFF3 file references a taxon outside that subset, ``bkbit`` downloads the full NCBI +taxonomy once and caches it in a per-user cache directory. To do that download up front +instead of mid-run - for example when building a container image or running air-gapped - +use the optional command: + +.. code-block:: bash + + $ bkbit download-ncbi-taxonomy + +Set ``BKBIT_DATA_DIR`` to control where the cache is stored, or ``BKBIT_NO_DOWNLOAD`` to +make an out-of-subset lookup raise an error instead of downloading. + +Command Line ............. ``bkbit gff2jsonld`` diff --git a/pyproject.toml b/pyproject.toml index a5e1b22..e6d4dfb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ authors = [ ] description = "A library for using brain-bican data models" readme = "README.md" -requires-python = ">=3.7" +requires-python = ">=3.9" classifiers = [ "Programming Language :: Python :: 3", "Operating System :: OS Independent", @@ -21,9 +21,19 @@ dependencies = [ "pandas", "click", "schemasheets", + "platformdirs", + "rdflib", + "requests", + "tqdm", ] dynamic = ["version"] +[tool.setuptools.packages.find] +include = ["bkbit*"] + +[tool.setuptools.package-data] +"bkbit.utils.ncbi_taxonomy_data" = ["*.json.gz"] + [project.scripts] bkbit = "bkbit.cli:cli"