code/add_structure.py

import re
import time
import json
import zlib
from xml.etree import ElementTree
from urllib.parse import urlparse, parse_qs, urlencode
import requests
import unipressed
from requests.adapters import HTTPAdapter, Retry
from unipressed import IdMappingClient
"""
## Code adapted from UniProt documentation.
def get_pdb_ids_2(protein_id):
    POLLING_INTERVAL = 5
    API_URL = "https://rest.uniprot.org"

    retries = Retry(total=5, backoff_factor=0.25, status_forcelist=[500, 502, 503, 504])
    session = requests.Session()
    session.mount("https://", HTTPAdapter(max_retries=retries))

    def check_response(response):
        try:
            response.raise_for_status()
        except requests.HTTPError:
            print(response.json())
            raise

    def submit_id_mapping(from_db, to_db, ids):
        request = requests.post(
            f"{API_URL}/idmapping/run",
            data={"from": from_db, "to": to_db, "ids": ids},
        )
        check_response(request)
        if check_response != None:
        	return request.json()["jobId"]
        else:
        	return None

    def get_next_link(headers):
        re_next_link = re.compile(r'<(.+)>; rel="next"')
        if "Link" in headers:
            match = re_next_link.match(headers["Link"])
            if match:
                return match.group(1)

    def check_id_mapping_results_ready(job_id):
        print('entered')
        while True:
            print('True')
            print('HR-1')
            try:
                request = session.get(f"{API_URL}/idmapping/status/{job_id}")
            except requests.exceptions.RetryError:
                print('eneted')
                request = None
            print('HR0-22')
            check_response(request)
            j = request.json()
            print('HR0')
            try:
                print('HR1')
                if "jobStatus" in j:
                    print('HR2')
                    if j["jobStatus"] == "RUNNING":
                        print(f"Retrying in {POLLING_INTERVAL}s")
                        time.sleep(POLLING_INTERVAL)
                    else:
                        print('HR3')
                        raise Exception(j["jobStatus"])
            except:
                print('HR4')
                requests.exceptions.RetryError
            else:
                print('HR4')
                return bool(j["results"] or j["failedIds"])

    def get_batch(batch_response, file_format, compressed):
        batch_url = get_next_link(batch_response.headers)
        while batch_url:
            batch_response = session.get(batch_url)
            batch_response.raise_for_status()
            yield decode_results(batch_response, file_format, compressed)
            batch_url = get_next_link(batch_response.headers)

    def combine_batches(all_results, batch_results, file_format):
        if file_format == "json":
            for key in ("results", "failedIds"):
                if key in batch_results and batch_results[key]:
                    all_results[key] += batch_results[key]
        elif file_format == "tsv":
            return all_results + batch_results[1:]
        else:
            return all_results + batch_results
        return all_results

    def get_id_mapping_results_link(job_id):
        url = f"{API_URL}/idmapping/details/{job_id}"

        request = session.get(url)
        check_response(request)
        return request.json()["redirectURL"]

    def decode_results(response, file_format, compressed):
        if compressed:
            decompressed = zlib.decompress(response.content, 16 + zlib.MAX_WBITS)
            if file_format == "json":
                j = json.loads(decompressed.decode("utf-8"))
                return j
            elif file_format == "tsv":
                return [line for line in decompressed.decode("utf-8").split("\n") if line]
            elif file_format == "xlsx":
                return [decompressed]
            elif file_format == "xml":
                return [decompressed.decode("utf-8")]
            else:
                return decompressed.decode("utf-8")
        elif file_format == "json":
            return response.json()
        elif file_format == "tsv":
            return [line for line in response.text.split("\n") if line]
        elif file_format == "xlsx":
            return [response.content]
        elif file_format == "xml":
            return [response.text]
        return response.text

    def get_xml_namespace(element):
        m = re.match(r"\{(.*)\}", element.tag)
        return m.groups()[0] if m else ""

    def merge_xml_results(xml_results):
        merged_root = ElementTree.fromstring(xml_results[0])
        for result in xml_results[1:]:
            root = ElementTree.fromstring(result)
            for child in root.findall("{http://uniprot.org/uniprot}entry"):
                merged_root.insert(-1, child)
        ElementTree.register_namespace("", get_xml_namespace(merged_root[0]))
        return ElementTree.tostring(merged_root, encoding="utf-8", xml_declaration=True)


    def get_id_mapping_results_search(url):
        parsed = urlparse(url)
        query = parse_qs(parsed.query)
        file_format = query["format"][0] if "format" in query else "json"
        if "size" in query:
            size = int(query["size"][0])
        else:
            size = 500
            query["size"] = size
        compressed = (
            query["compressed"][0].lower() == "true" if "compressed" in query else False
        )
        parsed = parsed._replace(query=urlencode(query, doseq=True))
        url = parsed.geturl()
        request = session.get(url)
        check_response(request)
        results = decode_results(request, file_format, compressed)
        total = int(request.headers["x-total-results"])
        for i, batch in enumerate(get_batch(request, file_format, compressed), 1):
            results = combine_batches(results, batch, file_format)
        if file_format == "xml":
            return merge_xml_results(results)
        return results


    job_id = submit_id_mapping(
        from_db="UniProtKB_AC-ID", to_db="PDB", ids=protein_id
    )
    print('skhfkh')
    print(submit_id_mapping(
        from_db="UniProtKB_AC-ID", to_db="PDB", ids=protein_id
    ))
    print('nor', check_id_mapping_results_ready(job_id))
    if check_id_mapping_results_ready(job_id):
        link = get_id_mapping_results_link(job_id)
        results = get_id_mapping_results_search(link)
        return [i['to'] for i in results['results']]
    else:
        print('no i am here')
        return None
def get_pdb_ids(protein_id):
    try:
        request = IdMappingClient.submit(
            source="UniProtKB_AC-ID", dest="PDB", ids={protein_id})

        try:
            pdb_list = list(request.each_result())
            time.sleep(1)
            return [i['to'] for i in pdb_list]
        except unipressed.id_mapping.core.IdMappingError:
            print('I AM HERE 1')
            get_pdb_ids_2(protein_id)
    except requests.exceptions.HTTPError:
        print('I AM HERE 2')
        get_pdb_ids_2(protein_id)
    except KeyError:
        print('I AM HERE 3')
        get_pdb_ids_2(protein_id)
"""

def get_pdb_ids(protein_id):
    try:
        request = IdMappingClient.submit(
            source="UniProtKB_AC-ID", dest="PDB", ids={protein_id})
        pdb_list = list(request.each_result())
        return [i['to'] for i in pdb_list]
    except requests.exceptions.HTTPError:
        return  []
    except unipressed.id_mapping.core.IdMappingError:
        print('IdMappingError caused by UniProt API service, please try later.')
        return  []
    except KeyError:
        return  []