FileIO/python/DQM.py

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