Belle II Software development
b2test_utils_analysis.py
1#!/usr/bin/env python3
2
3
10
11import os
12import sys
13import subprocess
14from contextlib import nullcontext
15import unittest
16import glob
17from basf2 import find_file
18from b2test_utils import configure_logging_for_tests, clean_working_directory
19
20
21class ExamplesTest(unittest.TestCase):
22 """Test to run all example scripts."""
23
24 @unittest.skipIf(not os.getenv('BELLE2_EXAMPLES_DATA_DIR'),
25 "$BELLE2_EXAMPLES_DATA_DIR not found.")
26 @unittest.skipIf(not os.getenv('BELLE2_VALIDATION_DATA_DIR'),
27 "$BELLE2_VALIDATION_DATA_DIR not found.")
28 def _test_examples_dir(self, path_to_glob, broken=None, filepattern="", nevents=10, cleanup=False):
29 """
30 Internal function to test a directory full of example scripts with an optional list of broken scripts to be skipped.
31
32 Parameters:
33 path_to_glob (str): the path to search for scripts
34 broken (list(str)): (optional) scripts that are known to be broken and can be skipped
35 """
36 if broken is None:
37 broken = []
38 configure_logging_for_tests()
39 all_egs = sorted(glob.glob(find_file(path_to_glob) + f"/{filepattern}*.py"))
40 for eg in all_egs:
41 with (clean_working_directory() if cleanup else nullcontext()):
42 filename = os.path.basename(eg)
43 if filename not in broken:
44 with self.subTest(msg=filename):
45 outputfilename = filename.replace('.py', '.root')
46 result = subprocess.run(['basf2', '-n', f'{nevents}', eg, '-o',
47 outputfilename], stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
48 if result.returncode != 0:
49 # failure running example so let's print the output
50 # on stderr so it's not split from output of unittest
51 # done like this since we don't want to decode/encode utf8
52 sys.stdout.buffer.write(result.stdout)
53 self.assertEqual(result.returncode, 0)
54
55 if os.path.exists(outputfilename) and isNtuple(outputfilename):
56 scanTTree(outputfilename)
57
58
59def scanTTree(filename):
60 from ROOT import TFile
61
62 tfile = TFile(filename, "READ")
63 print(f"TFile: {filename}")
64
65 # get lists of TTree in the filename
66 ttrees = [key.GetName() for key in tfile.GetListOfKeys() if key.GetClassName() == "TTree"]
67
68 for ttree_name in ttrees:
69 if ttree_name == "persistent":
70 continue
71 print(f"TTree: {ttree_name}")
72
73 # get TTree object
74 ttree = tfile.Get(ttree_name)
75 num_entries = ttree.GetEntries()
76
77 if num_entries == 0:
78 print("No entry found")
79 # print name all TBranches in the TTree
80 for branch in ttree.GetListOfBranches():
81 branch_name = branch.GetName()
82 print(f"TBranch: {branch_name}")
83
84 else:
85 ttree.GetEntry(0)
86 # print name and value of all TBranches in the TTree
87 for branch in ttree.GetListOfBranches():
88 branch_name = branch.GetName()
89 branch_value = getattr(ttree, branch_name)
90
91 if isinstance(branch_value, float):
92 print(f"TBranch: {branch_name}, {branch_value:.4g}")
93 else:
94 print(f"TBranch: {branch_name}, {branch_value}")
95
96 tfile.Close()
97
98
99def isNtuple(filename):
100 from ROOT import Belle2, TFile
101
102 tfile = TFile(filename, "READ")
103 ttree = tfile.Get("persistent")
104 if not ttree:
105 return False
106 ttree.GetEntry(0)
107
108 metadata = Belle2.FileMetaData(ttree.FileMetaData)
109
110 return metadata.getDataDescription()["isNtupleMetaData"]
Metadata information about a file.
_test_examples_dir(self, path_to_glob, broken=None, filepattern="", nevents=10, cleanup=False)