Changed to use enums rather than strings

Author: James Parkhurst

profet/alphafold.py

@@ -9,6 +9,7 @@
 from requests_html import HTMLSession
 import requests
 from bs4 import BeautifulSoup
+from .enums import FileType
 
 
 class Alphafold_DB:
@@ -40,11 +41,13 @@ def check_structure(self, uniprot_id: str) -> bool:
 
         """
         uniprot_id = uniprot_id.upper()
-        url = self.make_url(uniprot_id, "pdb")
+        url = self.make_url(uniprot_id, FileType.pdb)
         r = requests.get(url)
         return r.status_code != 404
 
-    def get_file_url(self, uniprot_id: str, filetype: str = "cif") -> str:
+    def get_file_url(
+        self, uniprot_id: str, filetype: FileType = FileType.cif
+    ) -> str:
         """
         Get file url relative to an id from the the Alphafold entry page
 
@@ -56,6 +59,7 @@ def get_file_url(self, uniprot_id: str, filetype: str = "cif") -> str:
             The URL of the file to download
 
         """
+        filetype = FileType(filetype)
 
         # Do we recognise the filetpye, otherwise raise an exception.
         if filetype in ["pdb", "cif"]:
@@ -77,7 +81,9 @@ def get_file_url(self, uniprot_id: str, filetype: str = "cif") -> str:
         # Return the URL
         return url["href"]
 
-    def make_url(self, uniprot_id: str, filetype: str = "cif") -> str:
+    def make_url(
+        self, uniprot_id: str, filetype: FileType = FileType.cif
+    ) -> str:
         """
         Make the URL for the protein
 
@@ -89,6 +95,7 @@ def make_url(self, uniprot_id: str, filetype: str = "cif") -> str:
             The URL of the file to download
 
         """
+        filetype = FileType(filetype)
 
         uniprot_id = uniprot_id.upper()
         af_id = "AF-" + uniprot_id + "-F1"
@@ -109,7 +116,7 @@ def make_url(self, uniprot_id: str, filetype: str = "cif") -> str:
     def get_pdb(
         self,
         uniprot_id: str,
-        filetype: str = "cif",
+        filetype: FileType = FileType.cif,
     ) -> tuple:
         """
         Returns pdb/cif as strings, saves to file if requested.
@@ -123,6 +130,7 @@ def get_pdb(
             Tuple containing the filename and file from the database
 
         """
+        filetype = FileType(filetype)
         # af_id = self.df.loc[self.df['Uniprot_ID'] == uniprot_id.upper()]["AF_ID"].to_numpy()[0]
         # version = self.df.loc[self.df['Uniprot_ID'] == uniprot_id.upper()]["version"].to_numpy()[0]
         # af_id = "AF-"+uniprot_id+"-F1"

profet/cache.py

@@ -1,5 +1,6 @@
 import json
 import os
+from .enums import FileType
 
 
 class PDBFileCache(object):
@@ -35,7 +36,7 @@ def __init__(self, directory: str = None):
         # The manifest filename
         self.manifest = os.path.join(self.directory, "manifest.txt")
 
-    def path(self, uniprot_id: str, filetype: str = "cif") -> str:
+    def path(self, uniprot_id: str, filetype: FileType = FileType.cif) -> str:
         """
         Get the proposed path
 
@@ -47,6 +48,7 @@ def path(self, uniprot_id: str, filetype: str = "cif") -> str:
             The absolute path
 
         """
+        filetype = FileType(filetype)
         assert filetype in ["pdb", "cif"]
         return os.path.join(self.directory, uniprot_id.lower()) + "." + filetype
 
@@ -64,7 +66,7 @@ def find(self, uniprot_id: str) -> list:
         return [
             filename
             for filename in [
-                self.path(uniprot_id, filetype) for filetype in ["pdb", "cif"]
+                self.path(uniprot_id, filetype) for filetype in FileType
             ]
             if os.path.exists(filename)
         ]
@@ -137,7 +139,11 @@ def items(self):
                 yield uniprot_id, self.path(uniprot_id, filetype[1:])
 
     def _update_manifest(
-        self, uniprot_id: str, fileorigin: str, filetype: str, filename: str
+        self,
+        uniprot_id: str,
+        fileorigin: str,
+        filetype: FileType,
+        filename: str,
     ):
         """
         Update the manifest file
@@ -149,6 +155,7 @@ def _update_manifest(
             filename: The filename
 
         """
+        filetype = FileType(filetype)
 
         # Read the current manifest
         if os.path.exists(self.manifest):

profet/enums.py

@@ -0,0 +1,6 @@
+from enum import Enum
+
+
+class FileType(str, Enum):
+    pdb = "pdb"
+    cif = "cif"

profet/pdb.py

@@ -1,5 +1,6 @@
 from urllib.request import urlretrieve
 from rcsbapi.search import TextQuery
+from .enums import FileType
 
 
 class PDB_DB:
@@ -50,7 +51,9 @@ def check_structure(self, uniprot_id: str) -> bool:
             return True
         return False
 
-    def make_url(self, uniprot_id: str, filetype: str = "pdb") -> str:
+    def make_url(
+        self, uniprot_id: str, filetype: FileType = FileType.pdb
+    ) -> str:
         """
         Make the URL for the protein
 
@@ -63,14 +66,15 @@ def make_url(self, uniprot_id: str, filetype: str = "pdb") -> str:
 
         """
 
+        filetype = FileType(filetype)
         uniprot_id = uniprot_id.upper()
         url = f"https://files.rcsb.org/download/{uniprot_id}.{filetype}"
         return url
 
     def get_pdb(
         self,
         uniprot_id: str,
-        filetype: str = "cif",
+        filetype: FileType = FileType.cif,
     ) -> tuple:
         """
         Returns pdb/cif as strings, saves to file if requested
@@ -84,6 +88,7 @@ def get_pdb(
 
         """
 
+        filetype = FileType(filetype)
         if not self.results:
             self.results = [self.uniprot_id_to_pdb_id(uniprot_id)]
 
@@ -97,9 +102,9 @@ def get_pdb(
                 filedata = file.read()
         except Exception:
             if filetype == "pdb":
-                filetype = "cif"
+                filetype = FileType.cif
             else:
-                filetype = "pdb"
+                filetype = FileType.pdb
             url = self.make_url(pdb_id, filetype)
             filename, result = urlretrieve(url)
             with open(filename) as file:

profet/profet.py

@@ -2,6 +2,7 @@
 from .pdb import PDB_DB
 from .cache import PDBFileCache
 from .cleaver import Cleaver
+from .enums import FileType
 import os
 
 
@@ -54,7 +55,7 @@ def cache(self) -> PDBFileCache:
     def file_from_db(
         self,
         prot_id: str,
-        filetype: str = "cif",
+        filetype: FileType = FileType.cif,
         db: str = "pdb",
     ) -> tuple:
         """
@@ -69,14 +70,15 @@ def file_from_db(
             Tuple containing the filename and file from the database
 
         """
+        filetype = FileType(filetype)
         return {"pdb": self.pdb.get_pdb, "alphafold": self.alpha.get_pdb}[db](
             prot_id, filetype=filetype
         )
 
     def get_file(
         self,
         uniprot_id: str,
-        filetype: str = "cif",
+        filetype: FileType = FileType.cif,
         filesave: bool = False,
         db: str = "pdb",
     ) -> tuple:
@@ -96,6 +98,7 @@ def get_file(
             2. File from the database, or None if it is not available in any database.
 
         """
+        filetype = FileType(filetype)
 
         # Get the PDB cache
         cache = PDBFileCache(directory=self.save_directory)
@@ -211,7 +214,7 @@ def remove(
         """
         # Get the PDB cache
         cache = PDBFileCache(directory=self.save_directory)
-        identifier, _, _ = self.pdb.get_pdb(uniprot_id, filetype="cif")
+        identifier, _, _ = self.pdb.get_pdb(uniprot_id, filetype=FileType.cif)
         if uniprot_id in cache:
             filename = cache[uniprot_id]
             signal_list = self.Cleaver.signal_residuenumbers_requester(