Source code for schrodinger.protein.gpcr.gpcrdb

"""
This module is for downloading sequence and residue data from the GPCR DB.

Copyright Schrodinger, LLC. All rights reserved.
"""

from json.decoder import JSONDecodeError
import time

import backoff
import networkx as nx
import requests


def _load_protein_families(data):
    """
    Load GPCR protein families into a directed acyclic graph. Node IDs are GPCR
    DB slugs.

    :param data: List of GPCR DB protein family data
    :type data: list[dict]

    :rtype: networkx.DiGraph
    """
    digraph = nx.DiGraph()
    for node in data:
        digraph.add_node(node['slug'], name=node['name'])
        parent = node['parent']
        if parent:
            digraph.add_edge(parent['slug'], node['slug'])
    return digraph


def _find_families(slug, digraph):
    """
    Recursive function to walk a directed acyclic graph of GPCR DB protein
    families to find all families that are ancestors of `slug`. Families are
    yielded from most specific (leaf) to most general (root).

    :param slug: GPCR DB slug (e.g. 001_001_001)
    :type slug: str

    :param digraph: Directed acyclic graph (returned from `_load_protein_families`)
    :type digraph: networkx.DiGraph

    :rtype: collections.abc.Iterable(str)
    """
    parent = next(digraph.predecessors(slug), None)
    if parent:
        yield digraph.nodes[slug]['name']
        yield from _find_families(parent, digraph)


def _find_families_from_graph(slug, digraph):
    return list(reversed(list(_find_families(slug, digraph))))


def find_families(slug, data):
    """
    Get all protein families that are ancestors of `slug`. Families are ordered
    from most general (root) to most specific (leaf).

    :param slug: GPCR DB slug (e.g. 001_001_001)
    :type slug: str

    :param data: List of GPCR DB protein family data
    :type data: list[dict]

    :rtype: list[str]
    """
    digraph = _load_protein_families(data)
    return _find_families_from_graph(slug, digraph)


def get_protein_families(force=False):
    """
    Download protein family data from GPCR DB.

    :rtype: list[dict]
    """
    url = 'https://gpcrdb.org/services/proteinfamily'
    response = requests.get(url)
    response.raise_for_status()
    return response.json()


def get_url(entry_name):
    """
    Get the GPCR DB URL for the given entry

    :param entry_name: GPCR DB entry name (e.g. 5ht1a_human)
    :type entry_name: str
    """
    return f"https://gpcrdb.org/protein/{entry_name}"


def download_all_entry_data():
    """
    Download data for all entries in the GPCR DB.
    """

    used_entries = set()
    saved_families = {}

    family_data = get_protein_families()
    family_graph = _load_protein_families(family_data)

    for one_family in family_data:
        slug = one_family['slug']
        yield from _get_family_data(slug,
                                    family_graph,
                                    used_entries=used_entries,
                                    saved_families=saved_families)


@backoff.on_exception(
    backoff.expo,
    (requests.exceptions.RequestException, JSONDecodeError),
    max_tries=5,
    jitter=backoff.full_jitter,
)
def _get_family_data(slug,
                     family_graph,
                     used_entries=None,
                     saved_families=None):
    """
    Download data for one family from the GPCR DB.

    :param slug: GPCR DB slug (e.g. 001_001_001)
    :type slug: str

    :param family_graph: Directed acyclic graph of protein families (returned
        from `_load_protein_families`)
    :type family_graph: networkx.DiGraph

    :param used_entries: Entries that have already been downloaded
    :type used_entries: set or NoneType

    :param saved_families: Dict of family ancestor lists keyed by family slug
    :type saved_families: dict or NoneType

    :return: tuples of data for each entry in the family
    :rtype: collections.abc.Iterable(tuple)
    """

    if used_entries is None:
        used_entries = set()

    if saved_families is None:
        saved_families = {}

    url = f'https://gpcrdb.org/services/proteinfamily/proteins/{slug}/'
    response = requests.get(url)
    response.raise_for_status()
    data = response.json()
    for one_entry in data:
        entry_name = one_entry['entry_name']
        if entry_name in used_entries:
            continue
        used_entries.add(entry_name)
        entry_slug = one_entry['family']
        if entry_slug not in saved_families:
            saved_families[entry_slug] = _find_families_from_graph(
                entry_slug, family_graph)
        families = saved_families[entry_slug]
        residues = _get_residue_annotations(entry_name)
        yield (entry_name, one_entry['residue_numbering_scheme'],
               one_entry['sequence'], families, residues)
        # TODO better rate limiting?
        time.sleep(1)


@backoff.on_exception(
    backoff.expo,
    (requests.exceptions.RequestException, JSONDecodeError),
    max_tries=5,
    jitter=backoff.full_jitter,
)
def _get_residue_annotations(entry_name):
    """
    Download residue data for one entry from the GPCR DB.

    :param entry_name: GPCR DB entry name (e.g. 5ht1a_human)
    :type entry_name: str
    """
    url = f'https://gpcrdb.org/services/residues/{entry_name}'
    response = requests.get(url)
    response.raise_for_status()
    return response.json()