added basic check for float aro numbers

BigDataBiology · Sep 10, 2023 · e5a5163 · e5a5163
1 parent b301535
commit e5a5163
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/argnorm/drug_categorization.py b/argnorm/drug_categorization.py
@@ -21,6 +21,9 @@ def get_immediate_drug_classes(aro_num: str) -> List[Tuple]:
             Each inner list contains the ARO number and name of the drug class in that order. [ARO:number, name].
     '''
 
+    if type(aro_num) == float or type(aro_num) == int:
+        aro_num = 'ARO:' + str(aro_num)
+
     gene = ARO[aro_num]
 
     confers_resistance_to_drug_class = any(r.name == 'confers_resistance_to_drug_class' for r in gene.relationships)