allow prefixing of identifiers

Author: JR-1991

pyenzyme/fetcher/pdb.py

@@ -194,6 +194,11 @@ def fetch_pdb(
     """
     # Get PDB data
     client = PDBClient()
+
+    # Allow prefixing with 'PDB:'
+    if pdb_id.lower().startswith("pdb:"):
+        pdb_id = pdb_id.split(":", 1)[-1]
+
     pdb_response = client.get_entry_by_id(pdb_id)
 
     if not pdb_response:

pyenzyme/fetcher/pubchem.py

@@ -145,7 +145,7 @@ def extract_by_preference(
 
 
 def fetch_pubchem(
-    cid: int,
+    cid: str,
     smallmol_id: Optional[str] = None,
     vessel_id: Optional[str] = None,
 ) -> v2.SmallMolecule:
@@ -163,9 +163,13 @@ def fetch_pubchem(
     Raises:
         ValueError: If the PubChem API request fails or required data is missing
     """
-    query = PubChemClient.from_cid(cid)
+    # Allow prefixing with 'CID:'
+    if cid.lower().startswith("pubchem:"):
+        cid = cid.split(":", 1)[-1]
+
+    query = PubChemClient.from_cid(int(cid))
     pc_compound = query.pc_compounds[0]
-    name = _extract_name(pc_compound, cid)
+    name = _extract_name(pc_compound, int(cid))
 
     if not smallmol_id:
         smallmol_id = process_id(name)

pyenzyme/fetcher/uniprot.py

@@ -129,6 +129,11 @@ def fetch_uniprot(
         ConnectionError: If the connection to the UniProt server fails
     """
     client = UniProtClient()
+
+    # Allow prefixing with 'uniprot:'
+    if uniprot_id.lower().startswith("uniprot:"):
+        uniprot_id = uniprot_id.split(":", 1)[-1]
+
     uniprot_entry = client.get_entry_by_id(uniprot_id)
 
     if not uniprot_entry:

tests/integration/test_composer.py

@@ -23,6 +23,25 @@ def test_compose(self):
         expected_doc = pe.read_enzymeml("tests/fixtures/compose/expected_compose.json")
         assert to_dict_wo_json_ld(doc) == to_dict_wo_json_ld(expected_doc)
 
+    def test_compose_with_prefix(self):
+        # Act
+        doc = pe.compose(
+            name="test",
+            vessel=pe.Vessel(
+                id="vessel",
+                name="vessel",
+                volume=1.0,
+                unit="ml",  # type: ignore
+            ),
+            proteins=["pdb:1A23"],
+            small_molecules=["CHEBI:32551"],
+            reactions=["RHEA:22864"],
+        )
+
+        # Assert
+        expected_doc = pe.read_enzymeml("tests/fixtures/compose/expected_compose.json")
+        assert to_dict_wo_json_ld(doc) == to_dict_wo_json_ld(expected_doc)
+
     def test_compose_no_vessel(self):
         # Act
         doc = pe.compose(

tests/integration/test_fetcher.py

@@ -57,6 +57,22 @@ def test_fetch_uniprot_to_protein(self):
         assert len(protein.references) == 1
         assert protein.references[0] == "https://www.uniprot.org/uniprotkb/P07327"
 
+    def test_fetch_uniprot_to_protein_with_prefix(self):
+        protein = fetch_uniprot("uniprot:P07327")
+        assert protein is not None
+
+        assert protein.id == "alcohol_dehydrogenase_1a"
+        assert protein.name == "Alcohol dehydrogenase 1A"
+        assert protein.constant is True
+        assert protein.ecnumber == "1.1.1.1"
+        assert protein.organism == "Homo sapiens"
+        assert protein.organism_tax_id == "9606"
+
+        assert protein.ld_id == "uniprot:P07327"
+        assert "uniprot:P07327" in protein.ld_type
+        assert len(protein.references) == 1
+        assert protein.references[0] == "https://www.uniprot.org/uniprotkb/P07327"
+
     def test_fetch_uniprot_to_protein_with_id(self):
         protein = fetch_uniprot("P07327", protein_id="p1")
         assert protein is not None
@@ -103,7 +119,19 @@ def test_fetch_rhea_to_reaction_invalid_id(self):
             fetch_rhea("INVALID_ID")
 
     def test_fetch_pubchem_to_small_molecule(self):
-        small_molecule = fetch_pubchem(cid=2244)
+        small_molecule = fetch_pubchem(cid="2244")
+        assert small_molecule is not None
+        assert small_molecule.id == "2_acetyloxybenzoic_acid"
+        assert small_molecule.name == "2-acetyloxybenzoic acid"
+        assert small_molecule.canonical_smiles == "CC(=O)OC1=CC=CC=C1C(=O)O"
+        assert small_molecule.inchikey == "BSYNRYMUTXBXSQ-UHFFFAOYSA-N"
+        assert (
+            small_molecule.inchi
+            == "InChI=1S/C9H8O4/c1-6(10)13-8-5-3-2-4-7(8)9(11)12/h2-5H,1H3,(H,11,12)"
+        )
+
+    def test_fetch_pubchem_to_small_molecule_with_prefix(self):
+        small_molecule = fetch_pubchem(cid="pubchem:2244")
         assert small_molecule is not None
         assert small_molecule.id == "2_acetyloxybenzoic_acid"
         assert small_molecule.name == "2-acetyloxybenzoic acid"
@@ -116,7 +144,7 @@ def test_fetch_pubchem_to_small_molecule(self):
 
     def test_fetch_pubchem_to_small_molecule_invalid_id(self):
         with pytest.raises(ValueError):
-            fetch_pubchem(cid=162176127617627)
+            fetch_pubchem(cid="162176127617627")
 
     def test_fetch_pdb_to_protein(self):
         protein = fetch_pdb("1a23")
@@ -131,6 +159,19 @@ def test_fetch_pdb_to_protein(self):
             == "AQYEDGKQYTTLEKPVAGAPQVLEFFSFFCPHCYQFEEVLHISDNVKKKLPEGVKMTKYHVNFMGGDLGKDLTQAWAVAMALGVEDKVTVPLFEGVQKTQTIRSASDIRDVFINAGIKGEEYDAAWNSFVVKSLVAQQEKAAADVQLRGVPAMFVNGKYQLNPQGMDTSNMDVFVQQYADTVKYLSEKK"
         )
 
+    def test_fetch_pdb_to_protein_with_prefix(self):
+        protein = fetch_pdb("pdb:1a23")
+        assert protein is not None
+        assert protein.id == "1a23_1"
+        assert (
+            protein.name
+            == "SOLUTION NMR STRUCTURE OF REDUCED DSBA FROM ESCHERICHIA COLI, MINIMIZED AVERAGE STRUCTURE"
+        )
+        assert (
+            protein.sequence
+            == "AQYEDGKQYTTLEKPVAGAPQVLEFFSFFCPHCYQFEEVLHISDNVKKKLPEGVKMTKYHVNFMGGDLGKDLTQAWAVAMALGVEDKVTVPLFEGVQKTQTIRSASDIRDVFINAGIKGEEYDAAWNSFVVKSLVAQQEKAAADVQLRGVPAMFVNGKYQLNPQGMDTSNMDVFVQQYADTVKYLSEKK"
+        )
+
     def test_fetch_pdb_to_protein_invalid_id(self):
         with pytest.raises(ValueError):
             fetch_pdb("INVALID_ID")