added test for column presence

MrOlm · Jan 18, 2023 · 498bf4b · 498bf4b
1 parent 08b9962
commit 498bf4b
Showing 1 changed file with 10 additions and 11 deletions.
diff --git a/drep/d_choose.py b/drep/d_choose.py
@@ -205,17 +205,16 @@ def score_genomes(genomes, Gdb, Edb=None, **kwargs):
         g2e = {}
 
     # Test for NaNs in Gdb['strain_heterogeneity']
-
-    if Gdb['strain_heterogeneity'].isnull().values.any():
-        kwargs['strain_heterogeneity_weight'] = 0
-        logging.warning ('NaNs found in Gdb strain_heterogeneity! '
-                         'This may have happened if CheckM results '
-                         'were imported for some but not all '
-                         'genomes. \n\n'
-                         'Setting strain_heterogeneity_weight to zero')
-        Gdb.drop('strain_heterogeneity',
-                 axis=1,
-                 inplace=True)
+    if 'strain_heterogeneity' in Gdb.columns:
+        if Gdb['strain_heterogeneity'].isnull().values.any():
+            kwargs['strain_heterogeneity_weight'] = 0
+            logging.warning ('NaNs found in Gdb strain_heterogeneity! '
+                             'This may have happened if CheckM results '
+                             'were imported for some but not all '
+                             'genomes. \n\n'
+                             'Setting strain_heterogeneity_weight to zero')
+            Gdb.drop('strain_heterogeneity',
+                     axis=1, inplace=True)
 
     Table = {'genome':[],'score':[]}
     for genome in genomes: