Spaces:
Runtime error
Runtime error
File size: 1,803 Bytes
fb9c306 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from dataclasses import dataclass
import requests
from fastapi import HTTPException
from graphgen.utils import logger
UNIPROT_BASE = "https://rest.uniprot.org/uniprotkb/search"
@dataclass
class UniProtSearch:
"""
UniProt Search client to search with UniProt.
1) Get the protein by accession number.
2) Search with keywords or protein names.
"""
def get_entry(self, accession: str) -> dict:
"""
Get the UniProt entry by accession number(e.g., P04637).
"""
url = f"{UNIPROT_BASE}/{accession}.json"
return self._safe_get(url).json()
def search(
self,
query: str,
*,
size: int = 10,
cursor: str = None,
fields: list[str] = None,
) -> dict:
"""
Search UniProt with a query string.
:param query: The search query.
:param size: The number of results to return.
:param cursor: The cursor for pagination.
:param fields: The fields to return in the response.
:return: A dictionary containing the search results.
"""
params = {
"query": query,
"size": size,
}
if cursor:
params["cursor"] = cursor
if fields:
params["fields"] = ",".join(fields)
url = UNIPROT_BASE
return self._safe_get(url, params=params).json()
@staticmethod
def _safe_get(url: str, params: dict = None) -> requests.Response:
r = requests.get(
url,
params=params,
headers={"Accept": "application/json"},
timeout=10,
)
if not r.ok:
logger.error("Search engine error: %s", r.text)
raise HTTPException(r.status_code, "Search engine error.")
return r
|