Ncbi Sequence Fetch

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""NCBI E-utilities wrapper for protein and nucleotide sequence retrieval.

Provides subcommands for the core NCBI E-utilities operations needed for
protein sequence retrieval: efetch, esearch, elink, CDS translation,
patent sequence search, and gene-to-protein lookup.

Rate-limited to 3 requests/second (10/s with NCBI_API_KEY).
"""

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "scienceskillscommon",
#   "python-dotenv",
# ]
# [tool.uv.sources]
# scienceskillscommon = { path = "../../scienceskillscommon" }
# ///

from __future__ import annotations

import argparse
import json
import os
import re
import sys
from typing import Any
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

import dotenv
from science_skills.scienceskillscommon import http_client

_CODON_TABLE = {
    'ATA': 'I',
    'ATC': 'I',
    'ATT': 'I',
    'ATG': 'M',
    'ACA': 'T',
    'ACC': 'T',
    'ACG': 'T',
    'ACT': 'T',
    'AAC': 'N',
    'AAT': 'N',
    'AAA': 'K',
    'AAG': 'K',
    'AGC': 'S',
    'AGT': 'S',
    'AGA': 'R',
    'AGG': 'R',
    'CTA': 'L',
    'CTC': 'L',
    'CTG': 'L',
    'CTT': 'L',
    'CCA': 'P',
    'CCC': 'P',
    'CCG': 'P',
    'CCT': 'P',
    'CAC': 'H',
    'CAT': 'H',
    'CAA': 'Q',
    'CAG': 'Q',
    'CGA': 'R',
    'CGC': 'R',
    'CGG': 'R',
    'CGT': 'R',
    'GTA': 'V',
    'GTC': 'V',
    'GTG': 'V',
    'GTT': 'V',
    'GCA': 'A',
    'GCC': 'A',
    'GCG': 'A',
    'GCT': 'A',
    'GAC': 'D',
    'GAT': 'D',
    'GAA': 'E',
    'GAG': 'E',
    'GGA': 'G',
    'GGC': 'G',
    'GGG': 'G',
    'GGT': 'G',
    'TCA': 'S',
    'TCC': 'S',
    'TCG': 'S',
    'TCT': 'S',
    'TTC': 'F',
    'TTT': 'F',
    'TTA': 'L',
    'TTG': 'L',
    'TAC': 'Y',
    'TAT': 'Y',
    'TAA': '*',
    'TAG': '*',
    'TGC': 'C',
    'TGT': 'C',
    'TGA': '*',
    'TGG': 'W',
}


EUTILS_BASE = 'https://eutils.ncbi.nlm.nih.gov/entrez/eutils'

_CLIENT = None


def get_api_client():
  """Returns the lazily initialized HttpClient."""
  global _CLIENT
  if _CLIENT is None:
    api_key = os.environ.get('NCBI_API_KEY', '')
    qps = 10.0 if api_key else 3.0
    _CLIENT = http_client.HttpClient(EUTILS_BASE, qps=qps)
  return _CLIENT


def _eutils_get(endpoint: str, params: dict[str, str | int]) -> str | None:
  """Sends a GET request to an NCBI E-utilities endpoint.

  Handles rate limiting, retries with backoff, and API key
  injection.

  Args:
    endpoint: E-utilities endpoint name (e.g. 'efetch.fcgi').
    params: Query parameters as a dict.

  Returns:
    Response text or None on failure.
  """
  api_key = os.environ.get('NCBI_API_KEY', '')
  if api_key:
    params['api_key'] = api_key

  query_string = urllib.parse.urlencode(params)
  full_url = f'{EUTILS_BASE}/{endpoint}'
  if query_string:
    full_url += f'?{query_string}'

  try:
    return get_api_client().fetch_text(full_url)
  except http_client.HttpError as e:
    print(f'{endpoint} error after all retires: {e}', file=sys.stderr)
    return None


def efetch(
    db: str,
    db_id: str | int,
    retmode: str = 'text',
    rettype: str = 'fasta',
) -> str | None:
  """Fetches data from NCBI efetch endpoint.

  Args:
    db: Database name (protein, nuccore, gene, pubmed, etc.)
    db_id: One or more comma-separated IDs.
    retmode: Return mode (text, xml, json).
    rettype: Return type (fasta, gb, gp, fasta_cds_aa, etc.)

  Returns:
    Response text or None on