Skip to content

Commit 5e0dd5b

Browse files
committed
[tutorials] Fixup staff.py tutorial
* Format code * Fix linter warnings * Close TFile explicitly (avoiding also an unused variable warning) * Close `.dat` file correctly using `with` statement * Replace invalid regular expression with just `split()`, which already splits on whitespaces already
1 parent a529540 commit 5e0dd5b

1 file changed

Lines changed: 44 additions & 37 deletions

File tree

tutorials/io/tree/staff.py

Lines changed: 44 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -12,25 +12,28 @@
1212
##
1313
## \author Wim Lavrijsen
1414

15-
import re, array, os
15+
import os
16+
1617
import ROOT
17-
from ROOT import TFile, TTree, gROOT, addressof
18+
from ROOT import TFile, TTree, addressof, gROOT
1819

1920
## A C/C++ structure is required, to allow memory based access
2021
gROOT.ProcessLine(
21-
"struct staff_t {\
22-
Int_t Category;\
23-
UInt_t Flag;\
24-
Int_t Age;\
25-
Int_t Service;\
26-
Int_t Children;\
27-
Int_t Grade;\
28-
Int_t Step;\
29-
Int_t Hrweek;\
30-
Int_t Cost;\
31-
Char_t Division[4];\
32-
Char_t Nation[3];\
33-
};" );
22+
"""struct staff_t {
23+
Int_t Category;
24+
UInt_t Flag;
25+
Int_t Age;
26+
Int_t Service;
27+
Int_t Children;
28+
Int_t Grade;
29+
Int_t Step;
30+
Int_t Hrweek;
31+
Int_t Cost;
32+
Char_t Division[4];
33+
Char_t Nation[3];
34+
};"""
35+
)
36+
3437

3538
## Function to read in data from ASCII file and fill the ROOT tree
3639
def staff():
@@ -40,33 +43,37 @@ def staff():
4043
# The input file cern.dat is a copy of the CERN staff data base
4144
# from 1988
4245

43-
f = TFile( 'staff.root', 'RECREATE' )
44-
tree = TTree( 'T', 'staff data from ascii file' )
45-
tree.Branch( 'staff', staff, 'Category/I:Flag:Age:Service:Children:Grade:Step:Hrweek:Cost' )
46-
tree.Branch( 'Divisions', addressof( staff, 'Division' ), 'Division/C' )
47-
tree.Branch( 'Nation', addressof( staff, 'Nation' ), 'Nation/C' )
46+
f = TFile("staff.root", "RECREATE")
47+
tree = TTree("T", "staff data from ascii file")
48+
tree.Branch("staff", staff, "Category/I:Flag:Age:Service:Children:Grade:Step:Hrweek:Cost")
49+
tree.Branch("Divisions", addressof(staff, "Division"), "Division/C")
50+
tree.Branch("Nation", addressof(staff, "Nation"), "Nation/C")
4851

4952
# note that the branches Division and Nation cannot be on the first branch
50-
fname = os.path.join(str(ROOT.gROOT.GetTutorialDir()), 'io', 'tree', 'cernstaff.dat')
51-
for line in open(fname).readlines():
52-
t = list(filter( lambda x: x, re.split( '\s+', line ) ) )
53-
staff.Category = int(t[0]) # assign as integers
54-
staff.Flag = int(t[1])
55-
staff.Age = int(t[2])
56-
staff.Service = int(t[3])
57-
staff.Children = int(t[4])
58-
staff.Grade = int(t[5])
59-
staff.Step = int(t[6])
60-
staff.Hrweek = int(t[7])
61-
staff.Cost = int(t[8])
62-
staff.Division = t[9] # assign as strings
63-
staff.Nation = t[10]
53+
fname = os.path.join(str(ROOT.gROOT.GetTutorialDir()), "io", "tree", "cernstaff.dat")
54+
with open(fname) as file:
55+
for line in file.readlines():
56+
t = line.split()
57+
staff.Category = int(t[0]) # assign as integers
58+
staff.Flag = int(t[1])
59+
staff.Age = int(t[2])
60+
staff.Service = int(t[3])
61+
staff.Children = int(t[4])
62+
staff.Grade = int(t[5])
63+
staff.Step = int(t[6])
64+
staff.Hrweek = int(t[7])
65+
staff.Cost = int(t[8])
66+
staff.Division = t[9] # assign as strings
67+
staff.Nation = t[10]
6468

65-
tree.Fill()
69+
tree.Fill()
6670

6771
tree.Print()
6872
tree.Write()
6973

74+
f.Close()
75+
76+
7077
#### run fill function if invoked on CLI
71-
if __name__ == '__main__':
72-
staff()
78+
if __name__ == "__main__":
79+
staff()

0 commit comments

Comments
 (0)