From 6e129b2492fbe44732fdc8e2015074276065dd85 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?V=C3=ADt=20Ku=C4=8Dera?= Date: Wed, 10 Apr 2024 12:39:48 +0200 Subject: [PATCH] Fix types --- PWGHF/D2H/Macros/hf_analysis_utils.py | 37 ++++++++++++++------------- 1 file changed, 19 insertions(+), 18 deletions(-) diff --git a/PWGHF/D2H/Macros/hf_analysis_utils.py b/PWGHF/D2H/Macros/hf_analysis_utils.py index e041c56f002..862dce80717 100644 --- a/PWGHF/D2H/Macros/hf_analysis_utils.py +++ b/PWGHF/D2H/Macros/hf_analysis_utils.py @@ -6,7 +6,14 @@ author: Fabrizio Grosa , CERN """ -import numpy as np +import numpy as np # pylint: disable=import-error + + +def make_list(object) -> list: + """ + Returns the object as a list if it is not a list already. + """ + return object if isinstance(object, list) else list(object) # pylint: disable=too-many-arguments @@ -63,7 +70,7 @@ def compute_fraction_fc( cross_sec_fd, raa_prompt=1.0, raa_fd=1.0, -): +) -> "tuple[list[float], list[float]]": """ Method to get fraction of prompt / FD fraction with fc method @@ -84,16 +91,13 @@ def compute_fraction_fc( - frac_fd: list of fraction of non-prompt D (central, min, max) """ - if not isinstance(cross_sec_prompt, list) and isinstance(cross_sec_prompt, float): - cross_sec_prompt = [cross_sec_prompt] - if not isinstance(cross_sec_fd, list) and isinstance(cross_sec_fd, float): - cross_sec_fd = [cross_sec_fd] - if not isinstance(raa_prompt, list) and isinstance(raa_prompt, float): - raa_prompt = [raa_prompt] - if not isinstance(raa_fd, list) and isinstance(raa_fd, float): - raa_fd = [raa_fd] + cross_sec_prompt = make_list(cross_sec_prompt) + cross_sec_fd = make_list(cross_sec_fd) + raa_prompt = make_list(raa_prompt) + raa_fd = make_list(raa_fd) - frac_prompt, frac_fd = [], [] + frac_prompt: list[float] = [] + frac_fd: list[float] = [] if acc_eff_prompt == 0: frac_fd_cent = 1.0 frac_prompt_cent = 0.0 @@ -141,7 +145,7 @@ def compute_fraction_nb( sigma_mb, raa_ratio=1.0, taa=1.0, -): +) -> "list[float]": """ Method to get fraction of prompt / FD fraction with Nb method @@ -165,13 +169,10 @@ def compute_fraction_nb( - frac: list of fraction of prompt (non-prompt) D (central, min, max) """ - if not isinstance(crosssection, list) and isinstance(crosssection, float): - crosssection = [crosssection] - - if not isinstance(raa_ratio, list) and isinstance(raa_ratio, float): - raa_ratio = [raa_ratio] + crosssection = make_list(crosssection) + raa_ratio = make_list(raa_ratio) - frac = [] + frac: list[float] = [] for i_sigma, sigma in enumerate(crosssection): for i_raa_ratio, raa_rat in enumerate(raa_ratio): raa_other = 1.0