ENH@get_KOs_db_X_ref:Squeezing by KO so orgs are in lists

dizak · Jul 19, 2018 · b706e3e · b706e3e
1 parent 82e3adf
commit b706e3e
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/prowler/apis.py b/prowler/apis.py
@@ -170,6 +170,7 @@ def get_KOs_db_X_ref(self,
                          target_db='genes',
                          skip_dwnld=False,
                          strip_prefix=True,
+                         squeeze=True,
                          sep="\t"):
         """
         Get desired KEGG's database entries linked with KEGG Orthology Group.
@@ -184,6 +185,9 @@ def get_KOs_db_X_ref(self,
             Key for KEGG_API.databases of desired database.
         skip_dwnld: bool, default <False>
             Read already downloaded file if <True>
+        squeeze: bool, default <True>
+            Compress list of the organisms to list for each KEGG Orthology
+            Group
 
         Uses
         -------
@@ -216,6 +220,10 @@ def get_KOs_db_X_ref(self,
             self.KOs_db_X_ref_df.replace({":.+": ""},
                                          regex=True,
                                          inplace=True)
+        if squeeze:
+            self.KOs_db_X_ref_df = self.KOs_db_X_ref_df.groupby(
+                by=[self.KEGG_ID]
+                )[self.ORG_GENE_ID].apply(list).to_frame().reset_index()
 
     def get_db_entries(self,
                        out_file_name):

diff --git a/tests/tests.py b/tests/tests.py
@@ -75,11 +75,15 @@ def setUpClass(cls):
                                            names=["ORF_ID",
                                                   "KEGG_ID"],
                                            dtype="object")
-        cls.ref_KOs_db_X_ref_df = pd.read_csv("test_data/ApisTests/ref_KOs_db_X_ref.csv",
-                                               sep="\t",
-                                               names=[apis.Columns.KEGG_ID,
-                                                      apis.Columns.ORG_GENE_ID],
-                                               index_col=[0])
+        cls.ref_KOs_db_X_ref_df = pd.read_csv(
+            "test_data/ApisTests/ref_KOs_db_X_ref.csv",
+            sep="\t",
+            names=[apis.Columns.KEGG_ID,
+                   apis.Columns.ORG_GENE_ID],
+            index_col=[0]
+            ).groupby(
+                by=[apis.Columns.KEGG_ID]
+                )[apis.Columns.ORG_GENE_ID].apply(list).to_frame().reset_index()
         cls.orgs_names = ["Haemophilus influenzae",
                           "Mycoplasma genitalium",
                           "Methanocaldococcus jannaschii",