Pdb Database

# 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.

"""Downloads PDB coordinate files (mmCIF or PDB) from RCSB."""

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


import argparse
import os
import sys

from science_skills.skills.scienceskillscommon import http_client

CLIENT = http_client.HttpClient("https://files-beta.rcsb.org", qps=5.0)


def sanitize_id(pdb_id: str) -> str:
  """Sanitize PDB ID."""
  pdb_id = pdb_id.lower().strip()
  if len(pdb_id) == 4:
    return f"pdb_0000{pdb_id}"
  return pdb_id


def download_files(args: argparse.Namespace):
  """Downloads coordinate files."""
  ids = [id.strip() for id in args.ids.split(",") if id.strip()]

  if len(ids) > 1_000:
    print(
        "Aborting: this script is not recommended for bulk downloads."
        "**Check with the user**, then consider downloading a copy of all"
        f" {args.format} files using the recommended bulk script at "
        "https://cdn.rcsb.org/wwpdb/docs/BetaArchiveBatchDownloader.py.",
        file=sys.stderr,
    )
    return

  # RCSB bulk download script suggests a higher QPS may be acceptable here (~5).
  estimated_time_secs = len(ids) / 5.0
  print(f"Estimated download time: {estimated_time_secs / 60:.1f} minutes).")

  if not os.path.exists(args.output_dir):
    os.makedirs(args.output_dir, exist_ok=True)

  for pdb_id in ids:
    sanitized_id = sanitize_id(pdb_id)
    shard_chars = sanitized_id[-3:-1]
    ext = "cif" if args.format == "mmcif" else "pdb"

    url = (
        f"/pub/wwpdb/pdb/data/entries/{shard_chars}/{sanitized_id}/"
        f"structures/{sanitized_id}.{ext}.gz"
    )

    output_path = os.path.join(args.output_dir, f"{sanitized_id}.{ext}.gz")

    print(f"Downloading {sanitized_id} from {url}...", file=sys.stderr)

    try:
      content = CLIENT.fetch_bytes(url)
      with open(output_path, "wb") as f:
        f.write(content)
      print(f"Saved to {output_path}", file=sys.stderr)
    except Exception as e:
      print(f"Failed to download {pdb_id}: {e}", file=sys.stderr)


def parse_args() -> argparse.Namespace:
  """Parse command line arguments."""
  parser = argparse.ArgumentParser(description="Download PDB coordinate files")
  parser.add_argument(
      "--format",
      type=str,
      required=True,
      choices=["mmcif", "pdb"],
      help="File format to download (mmcif or pdb)",
  )
  parser.add_argument(
      "--ids",
      type=str,
      required=True,
      help="Comma-separated list of PDB IDs",
  )
  parser.add_argument(
      "--output_dir",
      type=str,
      required=True,
      help="Directory to save files to",
  )
  return parser.parse_args()


if __name__ == "__main__":
  main_args = parse_args()
  download_files(main_args)

# 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.

"""Gets PDB data using the RCSB Data API (GraphQL)."""

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

import argparse
import sys
import urllib.parse

from science_skills.skills.scienceskillscommon import http_client

CLIENT = http_client.HttpClient("https://data.rcsb.org", qps=2.0)


def get_pdb_metadata(args: argparse.Namespace):
  """Executes a GraphQL query against the PDB Data API.

  Args:
    args: parsed command line arguments containing the query.
  """
  encoded_query = urllib.parse.quote(args.query.strip())
  url = f"https://data.rcsb.org/graphql?query={encoded_query}"
  print(f"Querying PDB Data API from {url}...", file=sys.stderr)

  content = CLIENT.fetch_bytes(url)
  with open(args.output, "w") as f:
    f.write(content.decode("utf-8"))


def parse_args() -> argparse.Namespace:
  """Parse command line arguments."""
  parser = argparse.ArgumentParser(
      description="Get PDB data using the RCSB Data API (GraphQL)"
  )
  parser.add_argument(
      "--query",
      type=str,
      required=True,
      help="GraphQL query string",
  )
  parser.add_argument(
      "--output",
      type=str,
      required=True,
      help="File to write the output to",
  )
  return parser.parse_args()


if __name__ == "__main__":
  main_args = parse_args()
  get_pdb_metadata(main_args)

# 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.

"""Fetches JSON PDB schemas and saves a greppable list of properties."""

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

import argparse
from typing import Any

from science_skills.skills.scienceskillscommon import http_client

SEARCH_CLIENT = http_client.HttpClient("https://search.rcsb.org", qps=2.0)
DATA_CLIENT = http_client.HttpClient("https://data.rcsb.org", qps=2.0)

# Mapping of API types to their corresponding schema URLs
API_URLS = {
    "search_structure": "https://search.rcsb.org/rcsbsearch/v2/metadata/schema",
    "search_chemical": (
        "https://search.rcsb.org/rcsbsearch/v2/metadata/chemical/schema"
    ),
    "data_entry": "https://data.rcsb.org/rest/v1/schema/entry",
    "data_polymer_entity": (
        "https://data.rcsb.org/rest/v1/schema/polymer_entity"
    ),
    "data_polymer_entity_instance": (
        "https://data.rcsb.org/rest/v1/schema/polymer_entity_instance"
    ),
    "data_assembly": "https://data.rcsb.org/rest/v1/schema/assembly",
    "data_non_polymer_entity": (
        "https://data.rcsb.org/rest/v1/schema/nonpolymer_entity"
    ),
    "data_non_polymer_entity_instance": (
        "https://data.rcsb.org/rest/v1/schema/nonpolymer_entity_instance"
    ),
    "data_branched_entity_instance": (
        "https://data.rcsb.org/rest/v1/schema/branched_entity_instance"
    ),
    "data_branched_entity": (
        "https://data.rcsb.org/rest/v1/schema/branched_entity"
    ),
    "data_chemical_component": "https://data.rcsb.org/rest/v1/schema/chem_comp",
}


def _collect_definitions(
    schema: dict[str, Any], prefix: str, config: argparse.Namespace
) -> set[str]:
  """Traverses a composite schema node and collects property definitions."""
  results = set()
  if "anyOf" in schema:
    for sub in schema["anyOf"]:
      results |= _collect_definitions(sub, prefix, config)
  if "oneOf" in schema:
    for sub in schema["oneOf"]:
      results |= _collect_definitions(sub, prefix, config)
  if "allOf" in schema:
    for sub in schema["allOf"]:
      results |= _collect_definitions(sub, prefix, config)

  type_val = schema.get("type")

  # Handle implicit object type if 'properties' is present
  if type_val == "object" or "properties" in schema:
    props = schema.get("properties", {})
    for key, val in props.items():
      full_name = f"{prefix}.{key}" if prefix else key
      results |= _process_property_node(val, full_name, config)

  # Handle implicit array type if 'items' is present
  elif type_val == "array" or "items" in schema:
    items = schema.get("items")
    if items:
      # Prefix doesn't include array_name.items, so no prefix update
      results |= _collect_definitions(items, prefix, config)

  return results


def _process_property_node(
    node: dict[str, Any], full_name: str, config: argparse.Namespace
) -> set[str]:
  """Processes an individual schema node and collects property definitions."""
  results = set()

  # Handle composite types first
  if "anyOf" in node:
    for sub in node["anyOf"]:
      results |= _process_property_node(sub, full_name, config)
  if "oneOf" in node:
    for sub in node["oneOf"]:
      results |= _process_property_node(sub, full_name, config)
  if "allOf" in node:
    for sub in node["allOf"]:
      results |= _process_property_node(sub, full_name, config)

  type_val = node.get("type")

  # Normalize type to list
  types = []
  if isinstance(type_val, list):
    types = type_val
  elif type_val:
    types = [type_val]

  is_object = "object" in types or "properties" in node
  is_array = "array" in types or "items" in node

  # Check if it has any primitive aspect
  # A property is primitive if it has a type that is NOT object or array
  has_primitive = False
  for t in types:
    if t not in ["object", "array"]:
      has_primitive = True

  # For arrays of primitives, metadata might be stored in 'items'
  meta_node = node
  if is_array and "items" in node and isinstance(node["items"], dict):
    if "rcsb_search_context" in node["items"] or "description" in node["items"]:
      meta_node = node["items"]

  # Output logic
  should_output = False
  if config.searchable_only:
    if "rcsb_search_context" in meta_node or "rcsb_search_context" in node:
      should_output = True
  else:
    if "description" in meta_node or has_primitive:
      should_output = True

  if should_output:
    description = (
        meta_node.get("description", node.get("description", ""))
        .replace("\n", " ")
        .strip()
    )

    # Truncate description if requested
    if (
        config.truncate_description > 0
        and len(description) > config.truncate_description
    ):
      description = description[: config.truncate_description] + "..."

    enum_values = meta_node.get("enum", node.get("enum"))
    enum_str = ""
    if enum_values:
      # Make a copy to avoid modifying original if we truncate
      vals = list(enum_values)
      if config.truncate_enums > 0 and len(vals) > config.truncate_enums:
        # Keep first N items and append "..."
        vals = vals[: config.truncate_enums] + ["..."]
      enum_str = f" {vals}"

    prop_type = node.get("type", "")
    if is_array and not prop_type:
      prop_type = "array"
    results.add(f"{full_name} ({prop_type}): {description}{enum_str}")

  # Recurse if it has object/array structure
  if is_object or is_array:
    results |= _collect_definitions(node, full_name, config)

  return results


def main():
  parser = argparse.ArgumentParser(
      description="Fetch JSON Schema and save a greppable list of properties."
  )
  parser.add_argument(
      "--api",
      type=str,
      choices=list(API_URLS.keys()),
      help="Api to use. Options: " + ",".join(list(API_URLS.keys())),
  )
  parser.add_argument(
      "--truncate-description",
      type=int,
      default=100,
      help=(
          "Number of characters to truncate description to (default: 100). Set"
          " to 0 to disable."
      ),
  )
  parser.add_argument(
      "--truncate-enums",
      type=int,
      default=0,
      help="Number of enum items to truncate to (default: 0 = no truncation).",
  )
  parser.add_argument(
      "--output",
      type=str,
      required=True,
      help="File to write the schema to.",
  )

  args = parser.parse_args()
  if "search" in args.api:
    args.searchable_only = True
  else:
    args.searchable_only = False

  if args.api:
    url = API_URLS[args.api]
    client = DATA_CLIENT if args.api.startswith("data_") else SEARCH_CLIENT
    print(f"Fetching schema from {url}...")
    try:
      data = client.fetch_json(url, timeout=30)
    except http_client.HttpError as e:
      print(f"Error fetching data from URL: {e}")
      return
  else:
    print("Error: No API type specified.")
    return

  # Start traversal
  results = _collect_definitions(schema=data, prefix="", config=args)

  # Sort and write
  sorted_lines = sorted(results)

  with open(args.output, "w") as f:
    for line in sorted_lines:
      f.write(line + "\n")

  print(
      f"Processed {len(sorted_lines)} properties. Output written to"
      f" {args.output}"
  )


if __name__ == "__main__":
  main()

# 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.

"""Searches PDB using the RCSB Search API v2.

This script allows executing structured queries against the PDB Search API
with tunable pagination, return types, and sorting.
"""

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

import argparse
import json
import sys
import urllib.parse

from science_skills.skills.scienceskillscommon import http_client

CLIENT = http_client.HttpClient("https://search.rcsb.org", qps=2.0)


def search_pdb(args: argparse.Namespace):
  """Executes a search against the PDB Search API.

  Args:
    args: parsed command line arguments containing the query and options.
  """
  try:
    parsed_query = json.loads(args.query)
  except json.JSONDecodeError as e:
    print(f"Error parsing --query as JSON: {e}", file=sys.stderr)
    sys.exit(1)

  if isinstance(parsed_query, dict) and "query" in parsed_query:
    # Payload already contains "query" key, it's a full request payload
    payload = parsed_query
  else:
    # Only "query" block provided
    payload = {"query": parsed_query}

  if args.return_type is not None:
    payload["return_type"] = args.return_type

  request_options = payload.get("request_options", {})

  if args.count_only:
    # Count-only mode: request 0 rows, just get total_count from response
    request_options["paginate"] = {"start": 0, "rows": 0}
    # Remove return_all_hits if present, since we don't want all results
    request_options.pop("return_all_hits", None)
  elif args.page_start is not None or args.rows is not None:
    # Remove return_all_hits so CLI pagination flags are not silently ignored.
    request_options.pop("return_all_hits", None)
    paginate = request_options.get("paginate", {})
    if args.page_start is not None:
      paginate["start"] = args.page_start
    if args.rows is not None:
      paginate["rows"] = args.rows
    request_options["paginate"] = paginate
  else:
    # Default behavior: return all hits if no pagination is specified
    request_options["return_all_hits"] = True

  if args.sort_by is not None:
    sort_item = {"sort_by": args.sort_by}
    if args.sort_direction is not None:
      sort_item["direction"] = args.sort_direction
    request_options["sort"] = [sort_item]

  if request_options:
    payload["request_options"] = request_options

  json_payload = json.dumps(payload, separators=(",", ":"))
  encoded_query = urllib.parse.quote(json_payload)
  url = f"https://search.rcsb.org/rcsbsearch/v2/query?json={encoded_query}"

  print(f"Querying PDB Search API from {url}...", file=sys.stderr)

  content = CLIENT.fetch_bytes(url)
  if args.count_only:
    # Parse the response to extract just the total count
    response_data = json.loads(content.decode("utf-8"))
    total_count = response_data.get("total_count", 0)
    count_result = {"total_count": total_count}
    print(f"Total count: {total_count}", file=sys.stderr)
    with open(args.output, "w") as f:
      json.dump(count_result, f, indent=2)
  else:
    with open(args.output, "w") as f:
      f.write(content.decode("utf-8"))


def parse_args() -> argparse.Namespace:
  """Parse command line arguments."""
  parser = argparse.ArgumentParser(
      description="Search PDB using the RCSB Search API v2"
  )
  parser.add_argument(
      "--query",
      type=str,
      required=True,
      help="JSON string of the query object or full request payload",
  )
  parser.add_argument(
      "--return_type",
      type=str,
      required=True,
      choices=[
          "entry",
          "assembly",
          "polymer_entity",
          "non_polymer_entity",
          "polymer_instance",
          "mol_definition",
      ],
      help=(
          "Type of returned object."
          "entry = [PDB ID]"
          "assembly = [PDB ID]-[ASSEMBLY ID]"
          "polymer_entity = [PDB ID]-[ENTITY ID]"
          "non_polymer_entity = [PDB ID]-[ENTITY ID]"
          "polymer_instance = [PDB ID].[LABEL ASYM ID]"
          "mol_definition = [CHEMICAL COMP ID] or [BIRD ID]"
      ),
  )
  parser.add_argument(
      "--sort_by",
      type=str,
      help=(
          "Attribute to sort by (commonly score or "
          "rcsb_accession_info.initial_release_date)"
      ),
  )
  parser.add_argument(
      "--sort_direction",
      type=str,
      choices=["asc", "desc"],
      help="Sort direction (used with --sort_by)",
  )
  parser.add_argument(
      "--page_start",
      type=int,
      help="Start index for pagination",
  )
  parser.add_argument(
      "--rows",
      type=int,
      help="Number of rows to return",
  )
  parser.add_argument(
      "--count-only",
      action="store_true",
      help=(
          "Return only the total count of matching entries, not the full"
          " result list. Useful when you need to know how many results match"
          " without downloading them all."
      ),
  )
  parser.add_argument(
      "--output",
      type=str,
      required=True,
      help="File to write the output to",
  )

  return parser.parse_args()


if __name__ == "__main__":
  main_args = parse_args()
  search_pdb(main_args)