Update code/add_structure.py

code/add_structure.py

@@ -7,7 +7,7 @@ from urllib.parse import urlparse, parse_qs, urlencode
 import requests
 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
@@ -165,5 +165,19 @@ def get_pdb_ids(protein_id):
         get_pdb_ids_2(protein_id)
     except KeyError:
         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  []