diff --git a/.gitignore b/.gitignore index 84bfc36c..7bc3c6ab 100644 --- a/.gitignore +++ b/.gitignore @@ -15,3 +15,4 @@ outputs menu_tools.egg-info dist ph2-menu-tools +prod_ph2-menu-tools diff --git a/menu_tools/utils/utils.py b/menu_tools/utils/utils.py index a0003b5d..6f9ff76f 100644 --- a/menu_tools/utils/utils.py +++ b/menu_tools/utils/utils.py @@ -58,19 +58,19 @@ def get_branches(ntuple_path: str, tree: str, obj: str): ntuple = glob.glob(ntuple_path)[0] with uproot.open(ntuple) as f: all_branches = f[tree].keys() - if "GenTree" in tree: - prefix = "Generator/" - elif "L1PhaseII" in tree: - prefix = "L1PhaseII/" - elif "Events" in tree: - prefix = "" - - ## nano - if tree == "Events": - obj_branches = [x.split("_")[-1] for x in all_branches if x.startswith(obj)] - ## no nano - else: - obj_branches = [x.removeprefix(prefix + obj) for x in all_branches if x.startswith(prefix+obj)] + if "GenTree" in tree: + prefix = "Generator/" + elif "L1PhaseII" in tree: + prefix = "L1PhaseII/" + elif "Events" in tree: + prefix = "" + + ## nano + if tree == "Events": + obj_branches = [x.split("_")[-1] for x in all_branches if x.startswith(obj)] + ## no nano + else: + obj_branches = [x.removeprefix(prefix + obj) for x in all_branches if x.startswith(prefix+obj)] return obj_branches