From e1bf3fe38264bca9f9599ae8835696d3d9793405 Mon Sep 17 00:00:00 2001 From: "Alexander V. Hopp" Date: Tue, 10 Dec 2024 13:26:17 +0100 Subject: [PATCH] Improve ugly if-else block --- baybe/utils/chemistry.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/baybe/utils/chemistry.py b/baybe/utils/chemistry.py index 86e5c80b5..7e7e93d03 100644 --- 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