fixed branch slimming (#102)

scikit-hep · Jul 3, 2024 · f5e127a · f5e127a
1 parent 1ca6704
commit f5e127a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 13 deletions.
diff --git a/src/hepconvert/copy_root.py b/src/hepconvert/copy_root.py
@@ -221,7 +221,7 @@ def copy_root(
             tqdm = _utils.check_tqdm()
             progress_bar = tqdm.tqdm(desc="Trees copied")
             progress_bar.reset(total=number_of_items)
-    for t in trees:
+    for t in trees:  # pylint: disable=too-many-nested-blocks
         tree = f[t]
         count_branches = get_counter_branches(tree)
         kb = filter_branches(tree, keep_branches, drop_branches, count_branches)
@@ -251,19 +251,12 @@ def copy_root(
                             )
                         }
                     )
-                for key in group:
-                    if key in kb:
-                        del chunk[key]
+                    for key in group:
+                        if key in kb:
+                            del chunk[key]
             if first:
                 first = False
-                if drop_branches:
-                    branch_types = {
-                        name: array.type
-                        for name, array in chunk.items()
-                        if name not in drop_branches
-                    }
-                else:
-                    branch_types = {name: array.type for name, array in chunk.items()}
+                branch_types = {name: array.type for name, array in chunk.items()}
                 of.mktree(
                     tree.name,
                     branch_types,

diff --git a/tests/test_copy_root.py b/tests/test_copy_root.py
@@ -66,13 +66,14 @@ def test_keep_branches(tmp_path):
     hepconvert.copy_root(
         Path(tmp_path) / "drop_branches.root",
         skhep_testdata.data_path("uproot-HZZ.root"),
-        drop_branches=["Jet_*", "MClepton_*"],
+        keep_branches="MClepton_*",
         counter_name=lambda counted: "N" + counted,
         force=True,
     )
     original = uproot.open(skhep_testdata.data_path("uproot-HZZ.root"))
 
     file = uproot.open(Path(tmp_path) / "drop_branches.root")
+    file["events"].show()
     for key in original["events"].keys():
         if key.startswith("MClepton_"):
             assert key in file["events"].keys()
@@ -81,6 +82,7 @@ def test_keep_branches(tmp_path):
             )
         else:
             assert key not in file["events"].keys()
+    file.close()
 
 
 def test_hepdata_example(tmp_path):