Improve ugly if-else block

emdgroup · Dec 13, 2024 · e1bf3fe · e1bf3fe
1 parent ad59eb0
commit e1bf3fe
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/baybe/utils/chemistry.py b/baybe/utils/chemistry.py
@@ -149,15 +149,14 @@ def smiles_to_fingerprint_features(
     )
     name = f"{encoding.name}_"
     prefix = prefix + "_" if prefix else ""
-    if all("fingerprint" in f for f in fingerprint_encoder.get_feature_names_out()):
-        col_names = [
-            prefix + name + f.split("fingerprint")[1]
-            for f in fingerprint_encoder.get_feature_names_out()
-        ]
-    else:
-        col_names = [
-            prefix + name + f for f in fingerprint_encoder.get_feature_names_out()
-        ]
+    no_descriptor_names = all(
+        "fingerprint" in f for f in fingerprint_encoder.get_feature_names_out()
+    )
+    suffixes = [
+        f.split("fingerprint")[1] if no_descriptor_names else f
+        for f in fingerprint_encoder.get_feature_names_out()
+    ]
+    col_names = [prefix + name + suffix for suffix in suffixes]
     df = pd.DataFrame(features, columns=col_names, dtype=DTypeFloatNumpy)
 
     return df