FileIO/python/DQM.py

0001 #!/bin/env python3
0002
0003 from builtins import range
0004 import ROOT as R
0005 import os, re
0006
0007 class DQMReader(object):
0008     """
0009     Reader for  DQM IO and DQM root files.
0010     """
0011
0012     #defined DQMIO types, index is important!
0013     DQMIO_TYPES = ["Ints","Floats","Strings",
0014          "TH1Fs","TH1Ss","TH1Ds",
0015          "TH2Fs", "TH2Ss", "TH2Ds",
0016          "TH3Fs", "TProfiles","TProfile2Ds", "kNIndicies"]
0017
0018     def __init__(self, input_filename):
0019         self._root_file = R.TFile.Open(input_filename)
0020
0021         ioTest = self._root_file.Get("Indices")
0022         if bool(ioTest):
0023             self.type = "DQMIO"
0024         else:
0025             self.type = "ROOT"
0026
0027     def read_objects(self):
0028         if (self.type == "DQMIO"):
0029             return self.read_objects_dqmio()
0030         else:
0031             return self.read_objects_root()
0032
0033     def read_objects_dqmio(self):
0034         indices = self._root_file.Get("Indices")
0035
0036         for y in range(indices.GetEntries()):
0037             indices.GetEntry(y)
0038             # print indices.Run, indices.Lumi, indices.Type
0039
0040             if indices.Type == 1000:
0041                 # nothing is stored here
0042                 # see https://github.com/cms-sw/cmssw/blob/8be445ac6fd9983d69156199d4d1fd3350f05d92/DQMServices/FwkIO/plugins/DQMRootOutputModule.cc#L437
0043                 continue
0044
0045             object_type = self.DQMIO_TYPES[indices.Type]
0046             t_tree = self._root_file.Get(object_type)
0047
0048             for i in range(indices.FirstIndex, indices.LastIndex + 1):
0049                 t_tree.GetEntry(i)
0050
0051                 fullname = str(t_tree.FullName)
0052                 yield (fullname, t_tree.Value, )
0053
0054     def read_objects_root(self):
0055         xml_re = re.compile(r"^<(.+)>(.+)=(.*)<\/\1>$")
0056         def parse_directory(di):
0057             directory = self._root_file.GetDirectory(di)
0058             for key in directory.GetListOfKeys():
0059                 entry = key.GetName()
0060                 rtype = key.GetClassName()
0061                 fullpath = "%s/%s" % (di, entry)
0062
0063                 if (rtype == "TDirectoryFile"):
0064                     for k, v in parse_directory(fullpath):
0065                         yield (k, v, )
0066                 else:
0067                     obj = self._root_file.Get(fullpath)
0068                     if obj:
0069                         yield (fullpath, obj, )
0070                     else:
0071                         # special case to parse the xml abomination
0072                         m = xml_re.search(entry)
0073                         if m:
0074                             name = m.group(1)
0075                             typecode = m.group(2)
0076                             value = m.group(3)
0077
0078                             fp = "%s/%s" % (di, name)
0079                             yield (fp, value, )
0080                         else:
0081                             raise Exception("Invalid xml:" + entry)
0082
0083
0084         path_fix = re.compile(r"^\/Run \d+")
0085         for fullname, obj in parse_directory(""):
0086             f = fullname.replace("/DQMData", "")
0087             f = f.replace("/Run summary", "")
0088             f = path_fix.sub(r"", f)
0089             if f[0] == "/":
0090                 f = f[1:]
0091
0092             yield f, obj
0093
0094     def close(self):
0095         self._root_file.Close()
0096
0097 if __name__ == '__main__':
0098     import argparse
0099
0100     parser = argparse.ArgumentParser()
0101     parser.add_argument("-i", "--input", help = "Input DQMIO ROOT file")
0102     args = parser.parse_args()
0103
0104     reader = DQMReader(args.input)
0105
0106     for (fn, v) in reader.read_objects():
0107         if (hasattr(v, "ClassName")):
0108             print(fn, v.ClassName())
0109         else:
0110             print(fn, type(v))
0111
0112     reader.close()