#!/usr/bin/env @python@ # ROOT command line tools module: cmdLineUtils # Author: Julien Ripoche # Mail: julien.ripoche@u-psud.fr # Date: 20/08/15 """Contain utils for ROOT command line tools""" ########## # Stream redirect functions # The original code of the these functions can be found here : # http://stackoverflow.com/questions/4675728/redirect-stdout-to-a-file-in-python/22434262#22434262 # Thanks J.F. Sebastian !! from contextlib import contextmanager import os import sys # Support both Python2 and Python3 at the same time if sys.version_info.major > 2 : _input = input else: _input = raw_input def fileno(file_or_fd): """ Look for 'fileno' attribute. """ fd = getattr(file_or_fd, 'fileno', lambda: file_or_fd)() if not isinstance(fd, int): raise ValueError("Expected a file (`.fileno()`) or a file descriptor") return fd @contextmanager def streamRedirected(source=sys.stdout, destination=os.devnull): """ Redirect the output from source to destination. """ stdout_fd = fileno(source) # copy stdout_fd before it is overwritten #NOTE: `copied` is inheritable on Windows when duplicating a standard stream with os.fdopen(os.dup(stdout_fd), 'wb') as copied: source.flush() # flush library buffers that dup2 knows nothing about try: os.dup2(fileno(destination), stdout_fd) # $ exec >&destination except ValueError: # filename with open(destination, 'wb') as destination_file: os.dup2(destination_file.fileno(), stdout_fd) # $ exec > destination try: yield source # allow code to be run with the redirected stream finally: # restore source to its previous value #NOTE: dup2 makes stdout_fd inheritable unconditionally source.flush() os.dup2(copied.fileno(), stdout_fd) # $ exec >&copied def stdoutRedirected(): """ Redirect the output from sys.stdout to os.devnull. """ return streamRedirected(sys.stdout, os.devnull) def stderrRedirected(): """ Redirect the output from sys.stderr to os.devnull. """ return streamRedirected(sys.stderr, os.devnull) # The end of streamRedirected functions ########## ########## # Imports ## # redirect output (escape characters during ROOT importation...) with stdoutRedirected(): import ROOT # Silence Davix warning (see ROOT-7577) ROOT.PyConfig.IgnoreCommandLineOptions = True ROOT.gROOT.GetVersion() import argparse import glob import fnmatch import logging LOG_FORMAT = '%(levelname)s: %(message)s' logging.basicConfig(format=LOG_FORMAT) # The end of imports ########## ########## # Different functions to get a parser of arguments and options def _getParser(theHelp, theEpilog): """ Get a commandline parser with the defaults of the commandline utils. """ return argparse.ArgumentParser(description=theHelp, formatter_class=argparse.RawDescriptionHelpFormatter, epilog = theEpilog) def getParserSingleFile(theHelp, theEpilog=""): """ Get a commandline parser with the defaults of the commandline utils and a source file or not. """ parser = _getParser(theHelp, theEpilog) parser.add_argument("FILE", nargs='?', help="Input file") return parser def getParserFile(theHelp, theEpilog=""): """ Get a commandline parser with the defaults of the commandline utils and a list of source files. """ parser = _getParser(theHelp, theEpilog) parser.add_argument("FILE", nargs='+', help="Input file") return parser def getParserSourceDest(theHelp, theEpilog=""): """ Get a commandline parser with the defaults of the commandline utils, a list of source files and a destination file. """ parser = _getParser(theHelp, theEpilog) parser.add_argument("SOURCE", nargs='+', help="Source file") parser.add_argument("DEST", help="Destination file") return parser # The end of get parser functions ########## ########## # Several utils @contextmanager def _setIgnoreLevel(level): originalLevel = ROOT.gErrorIgnoreLevel ROOT.gErrorIgnoreLevel = level yield ROOT.gErrorIgnoreLevel = originalLevel def changeDirectory(rootFile,pathSplit): """ Change the current directory (ROOT.gDirectory) by the corresponding (rootFile,pathSplit) """ rootFile.cd() for directoryName in pathSplit: theDir = ROOT.gDirectory.Get(directoryName) if not theDir: logging.warning("Directory %s does not exist." %directoryName) return 1 else: theDir.cd() return 0 def createDirectory(rootFile,pathSplit): """ Add a directory named 'pathSplit[-1]' in (rootFile,pathSplit[:-1]) """ retcode = changeDirectory(rootFile,pathSplit[:-1]) if retcode == 0: ROOT.gDirectory.mkdir(pathSplit[-1]) return retcode def getFromDirectory(objName): """ Get the object objName from the current directory """ return ROOT.gDirectory.Get(objName) def isExisting(rootFile,pathSplit): """ Return True if the object, corresponding to (rootFile,pathSplit), exits """ changeDirectory(rootFile,pathSplit[:-1]) return ROOT.gDirectory.GetListOfKeys().Contains(pathSplit[-1]) def isDirectoryKey(key): """ Return True if the object, corresponding to the key, inherits from TDirectory """ classname = key.GetClassName() cl = ROOT.gROOT.GetClass(classname) return cl.InheritsFrom(ROOT.TDirectory.Class()) def isTreeKey(key): """ Return True if the object, corresponding to the key, inherits from TTree """ classname = key.GetClassName() cl = ROOT.gROOT.GetClass(classname) return cl.InheritsFrom(ROOT.TTree.Class()) def isTHnSparseKey(key): """ Return True if the object, corresponding to the key, inherits from THnSparse """ classname = key.GetClassName() cl = ROOT.gROOT.GetClass(classname) return cl.InheritsFrom(ROOT.THnSparse.Class()) def getKey(rootFile,pathSplit): """ Get the key of the corresponding object (rootFile,pathSplit) """ changeDirectory(rootFile,pathSplit[:-1]) return ROOT.gDirectory.GetKey(pathSplit[-1]) def isDirectory(rootFile,pathSplit): """ Return True if the object, corresponding to (rootFile,pathSplit), inherits from TDirectory """ if pathSplit == []: return True # the object is the rootFile itself else: return isDirectoryKey(getKey(rootFile,pathSplit)) def isTree(rootFile,pathSplit): """ Return True if the object, corresponding to (rootFile,pathSplit), inherits from TTree """ if pathSplit == []: return False # the object is the rootFile itself else: return isTreeKey(getKey(rootFile,pathSplit)) def getKeyList(rootFile,pathSplit): """ Get the list of keys of the directory (rootFile,pathSplit), if (rootFile,pathSplit) is not a directory then get the key in a list """ if isDirectory(rootFile,pathSplit): changeDirectory(rootFile,pathSplit) return ROOT.gDirectory.GetListOfKeys() else: return [getKey(rootFile,pathSplit)] def keyListSort(keyList): """ Sort list of keys by their names ignoring the case """ keyList.sort(key=lambda x: x.GetName().lower()) def tupleListSort(tupleList): """ Sort list of tuples by their first elements ignoring the case """ tupleList.sort(key=lambda x: x[0].lower()) def dirListSort(dirList): """ Sort list of directories by their names ignoring the case """ dirList.sort(key=lambda x: [n.lower() for n in x]) def keyClassSpliter(rootFile,pathSplitList): """ Return a list of directories and a list of keys corresponding to the other objects, for rootLs and rooprint use """ keyList = [] dirList = [] for pathSplit in pathSplitList: if pathSplit == []: dirList.append(pathSplit) elif isDirectory(rootFile,pathSplit): dirList.append(pathSplit) else: keyList.append(getKey(rootFile,pathSplit)) keyListSort(keyList) dirListSort(dirList) return keyList,dirList def openROOTFile(fileName, mode="read"): """ Open the ROOT file corresponding to fileName in the corresponding mode, redirecting the output not to see missing dictionnaries Returns: theFile (TFile) """ #with stderrRedirected(): with _setIgnoreLevel(ROOT.kError): theFile = ROOT.TFile.Open(fileName, mode) if not theFile: logging.warning("File %s does not exist", fileName) return theFile def openROOTFileCompress(fileName, compress, recreate): """ Open a ROOT file (like openROOTFile) with the possibility to change compression settings """ if compress != None and os.path.isfile(fileName): logging.warning("can't change compression settings on existing file") return None mode = "recreate" if recreate else "update" theFile = openROOTFile(fileName, mode) if compress != None: theFile.SetCompressionSettings(compress) return theFile def joinPathSplit(pathSplit): """ Join the pathSplit with '/' """ return "/".join(pathSplit) MANY_OCCURENCE_WARNING = "Several versions of '{0}' are present in '{1}'. Only the most recent will be considered." def manyOccurenceRemove(pathSplitList,fileName): """ Search for double occurence of the same pathSplit and remove them """ if len(pathSplitList) > 1: for n in pathSplitList: if pathSplitList.count(n) != 1: logging.warning(MANY_OCCURENCE_WARNING.format(joinPathSplit(n),fileName)) while n in pathSplitList and pathSplitList.count(n) != 1 : pathSplitList.remove(n) def patternToPathSplitList(fileName,pattern): """ Get the list of pathSplit of objects in the ROOT file corresponding to fileName that match with the pattern """ # Open ROOT file rootFile = openROOTFile(fileName) if not rootFile: return [] # Split pattern avoiding multiple slash problem patternSplit = [n for n in pattern.split("/") if n != ""] # Main loop pathSplitList = [[]] for patternPiece in patternSplit: newPathSplitList = [] for pathSplit in pathSplitList: if isDirectory(rootFile,pathSplit): changeDirectory(rootFile,pathSplit) newPathSplitList.extend( \ [pathSplit + [key.GetName()] \ for key in ROOT.gDirectory.GetListOfKeys() \ if fnmatch.fnmatch(key.GetName(),patternPiece)]) pathSplitList = newPathSplitList # No match if pathSplitList == []: logging.warning("can't find {0} in {1}".format(pattern,fileName)) # Same match (remove double occurrences from the list) manyOccurenceRemove(pathSplitList,fileName) return pathSplitList def fileNameListMatch(filePattern,wildcards): """ Get the list of fileName that match with objPattern """ if wildcards: return [os.path.expandvars(os.path.expanduser(i)) for i in glob.iglob(filePattern)] else: return [os.path.expandvars(os.path.expanduser(filePattern))] def pathSplitListMatch(fileName,objPattern,wildcards): """ Get the list of pathSplit that match with objPattern """ if wildcards: return patternToPathSplitList(fileName,objPattern) else: return [[n for n in objPattern.split("/") if n != ""]] def patternToFileNameAndPathSplitList(pattern,wildcards = True): """ Get the list of tuple containing both : - ROOT file name - list of splited path (in the corresponding file) of objects that matche Use unix wildcards by default """ rootFilePattern = "*.root" rootObjPattern = rootFilePattern+":*" httpRootFilePattern = "htt*://*.root" httpRootObjPattern = httpRootFilePattern+":*" xrootdRootFilePattern = "root://*.root" xrootdRootObjPattern = xrootdRootFilePattern+":*" s3RootFilePattern = "s3://*.root" s3RootObjPattern = s3RootFilePattern+":*" gsRootFilePattern = "gs://*.root" gsRootObjPattern = gsRootFilePattern+":*" pcmFilePattern = "*.pcm" pcmObjPattern = pcmFilePattern+":*" if fnmatch.fnmatch(pattern,httpRootObjPattern) or \ fnmatch.fnmatch(pattern,xrootdRootObjPattern) or \ fnmatch.fnmatch(pattern,s3RootObjPattern) or \ fnmatch.fnmatch(pattern,gsRootObjPattern): patternSplit = pattern.rsplit(":", 1) fileName = patternSplit[0] objPattern = patternSplit[1] pathSplitList = pathSplitListMatch(fileName,objPattern,wildcards) return [(fileName,pathSplitList)] if fnmatch.fnmatch(pattern,httpRootFilePattern) or \ fnmatch.fnmatch(pattern,xrootdRootFilePattern) or \ fnmatch.fnmatch(pattern,s3RootFilePattern) or \ fnmatch.fnmatch(pattern,gsRootFilePattern): fileName = pattern pathSplitList = [[]] return [(fileName,pathSplitList)] if fnmatch.fnmatch(pattern,rootObjPattern) or \ fnmatch.fnmatch(pattern,pcmObjPattern): patternSplit = pattern.split(":") filePattern = patternSplit[0] objPattern = patternSplit[1] fileNameList = fileNameListMatch(filePattern,wildcards) return [(fileName,pathSplitListMatch(fileName,objPattern,wildcards)) for fileName in fileNameList] if fnmatch.fnmatch(pattern,rootFilePattern) or \ fnmatch.fnmatch(pattern,pcmFilePattern): filePattern = pattern fileNameList = fileNameListMatch(filePattern,wildcards) pathSplitList = [[]] return [(fileName,pathSplitList) for fileName in fileNameList] logging.warning("{0}: No such file (or extension not supported)".format(pattern)) return [] # End of utils ########## ########## # Set of functions to put the arguments in shape def getArgs(parser): """ Get arguments corresponding to parser. """ return parser.parse_args() def getSourceListArgs(parser, wildcards = True): """ Create a list of tuples that contain source ROOT file names and lists of path in these files as well as the original arguments """ args = getArgs(parser) inputFiles = [] try: inputFiles = args.FILE except: inputFiles = args.SOURCE sourceList = \ [tup for pattern in inputFiles \ for tup in patternToFileNameAndPathSplitList(pattern,wildcards)] return sourceList, args def getSourceListOptDict(parser, wildcards = True): """ Get the list of tuples and the dictionary with options returns: sourceList: a list of tuples with one list element per file the first tuple entry being the root file, the second a list of subdirectories, each being represented as a list itself with a string per level e.g. rootls tutorial/tmva/TMVA.root:Method_BDT/BDT turns into [('tutorials/tmva/TMVA.root', [['Method_BDT','BDT']])] vars(args): a dictionary of matched options, e.g. {'longListing': False, 'oneColumn': False, 'treeListing': False, 'FILE': ['tutorials/tmva/TMVA.root:Method_BDT/BDT'] } """ sourceList, args = getSourceListArgs(parser, wildcards) if sourceList == []: logging.error("Input file(s) not found!") return sourceList, vars(args) def getSourceDestListOptDict(parser, wildcards = True): """ Get the list of tuples of sources, create destination name, destination pathSplit and the dictionary with options """ sourceList, args = getSourceListArgs(parser, wildcards) destList = \ patternToFileNameAndPathSplitList( \ args.DEST,wildcards=False) if destList != []: destFileName,destPathSplitList = destList[0] destPathSplit = destPathSplitList[0] else: destFileName = "" destPathSplit = [] return sourceList, destFileName, destPathSplit, vars(args) # The end of the set of functions to put the arguments in shape ########## ########## # Several functions shared by rootcp, rootmv and rootrm TARGET_ERROR = "target '{0}' is not a directory" OMITTING_ERROR = "omitting {0} '{1}'. Did you forget to specify the -r option for a recursive copy?" OVERWRITE_ERROR = "cannot overwrite non-directory '{0}' with directory '{1}'" def copyRootObject(sourceFile,sourcePathSplit,destFile,destPathSplit,oneSource,recursive,replace): """ Initialize the recursive function 'copyRootObjectRecursive', written to be as unix-like as possible """ retcode = 0 isMultipleInput = not (oneSource and sourcePathSplit != []) recursiveOption = recursive # Multiple input and un-existing or non-directory destination # TARGET_ERROR if isMultipleInput and destPathSplit != [] \ and not (isExisting(destFile,destPathSplit) \ and isDirectory(destFile,destPathSplit)): logging.warning(TARGET_ERROR.format(destPathSplit[-1])) retcode += 1 # Entire ROOT file or directory in input omitting "-r" option # OMITTING_FILE_ERROR or OMITTING_DIRECTORY_ERROR if not recursiveOption: if sourcePathSplit == []: logging.warning(OMITTING_ERROR.format( \ "file", sourceFile.GetName())) retcode += 1 elif isDirectory(sourceFile,sourcePathSplit): logging.warning(OMITTING_DIRECTORY_ERROR.format( \ "directory", sourcePathSplit[-1])) retcode += 1 # Run copyRootObjectRecursive function with the wish # to follow the unix copy behaviour if sourcePathSplit == []: retcode += copyRootObjectRecursive(sourceFile,sourcePathSplit, \ destFile,destPathSplit,replace) else: setName = "" if not isMultipleInput and (destPathSplit != [] \ and not isExisting(destFile,destPathSplit)): setName = destPathSplit[-1] objectName = sourcePathSplit[-1] if isDirectory(sourceFile,sourcePathSplit): if setName != "": createDirectory(destFile,destPathSplit[:-1]+[setName]) retcode += copyRootObjectRecursive(sourceFile,sourcePathSplit, \ destFile,destPathSplit[:-1]+[setName],replace) elif isDirectory(destFile,destPathSplit): if not isExisting(destFile,destPathSplit+[objectName]): createDirectory(destFile,destPathSplit+[objectName]) if isDirectory(destFile,destPathSplit+[objectName]): retcode += copyRootObjectRecursive(sourceFile,sourcePathSplit, \ destFile,destPathSplit+[objectName],replace) else: logging.warning(OVERWRITE_ERROR.format( \ objectName,objectName)) retcode += 1 else: logging.warning(OVERWRITE_ERROR.format( \ destPathSplit[-1],objectName)) retcode += 1 else: if setName != "": retcode += copyRootObjectRecursive(sourceFile,sourcePathSplit, \ destFile,destPathSplit[:-1],replace,setName) elif isDirectory(destFile,destPathSplit): retcode += copyRootObjectRecursive(sourceFile,sourcePathSplit, \ destFile,destPathSplit,replace) else: setName = destPathSplit[-1] retcode += copyRootObjectRecursive(sourceFile,sourcePathSplit, \ destFile,destPathSplit[:-1],replace,setName) return retcode DELETE_ERROR = "object {0} was not existing, so it is not deleted" def deleteObject(rootFile,pathSplit): """ Delete the object 'pathSplit[-1]' from (rootFile,pathSplit[:-1]) """ retcode = changeDirectory(rootFile,pathSplit[:-1]) if retcode == 0: fileName = pathSplit[-1] if isExisting(rootFile,pathSplit): ROOT.gDirectory.Delete(fileName+";*") else: logging.warning(DELETE_ERROR.format(fileName)) retcode += 1 return retcode def copyRootObjectRecursive(sourceFile,sourcePathSplit,destFile,destPathSplit,replace,setName=""): """ Copy objects from a file or directory (sourceFile,sourcePathSplit) to an other file or directory (destFile,destPathSplit) - Has the will to be unix-like - that's a recursive function - Python adaptation of a root input/output tutorial : $ROOTSYS/tutorials/io/copyFiles.C """ retcode = 0 replaceOption = replace seen = {} for key in getKeyList(sourceFile,sourcePathSplit): objectName = key.GetName() # write keys only if the cycle is higher than before if objectName not in seen.keys(): seen[objectName] = key else: if seen[objectName].GetCycle() < key.GetCycle(): seen[objectName] = key else: continue if isDirectoryKey(key): if not isExisting(destFile,destPathSplit+[objectName]): createDirectory(destFile,destPathSplit+[objectName]) if isDirectory(destFile,destPathSplit+[objectName]): retcode +=copyRootObjectRecursive(sourceFile, \ sourcePathSplit+[objectName], \ destFile,destPathSplit+[objectName],replace) else: logging.warning(OVERWRITE_ERROR.format( \ objectName,objectName)) retcode += 1 elif isTreeKey(key): T = key.GetMotherDir().Get(objectName+";"+str(key.GetCycle())) if replaceOption and isExisting(destFile,destPathSplit+[T.GetName()]): retcodeTemp = deleteObject(destFile,destPathSplit+[T.GetName()]) if retcodeTemp: retcode += retcodeTemp continue changeDirectory(destFile,destPathSplit) newT = T.CloneTree(-1,"fast") if setName != "": newT.SetName(setName) newT.Write() else: obj = key.ReadObj() if replaceOption and isExisting(destFile,destPathSplit+[setName]): changeDirectory(destFile,destPathSplit) otherObj = getFromDirectory(setName) if not otherObj == obj: retcodeTemp = deleteObject(destFile,destPathSplit+[setName]) if retcodeTemp: retcode += retcodeTemp continue else: obj.SetName(setName) changeDirectory(destFile,destPathSplit) obj.Write() else: obj.SetName(setName) changeDirectory(destFile,destPathSplit) obj.Write() elif issubclass(obj.__class__, ROOT.TCollection): # probably the object was written with kSingleKey changeDirectory(destFile,destPathSplit) obj.Write(setName, ROOT.TObject.kSingleKey) else: if setName != "": obj.SetName(setName) else: obj.SetName(objectName) changeDirectory(destFile,destPathSplit) obj.Write() obj.Delete() changeDirectory(destFile,destPathSplit) ROOT.gDirectory.SaveSelf(ROOT.kTRUE) return retcode FILE_REMOVE_ERROR = "cannot remove '{0}': Is a ROOT file" DIRECTORY_REMOVE_ERROR = "cannot remove '{0}': Is a directory" ASK_FILE_REMOVE = "remove '{0}' ? (y/n) : " ASK_OBJECT_REMOVE = "remove '{0}' from '{1}' ? (y/n) : " def deleteRootObject(rootFile, pathSplit, interactive, recursive): """ Remove the object (rootFile,pathSplit) -interactive : prompt before every removal -recursive : allow directory, and ROOT file, removal """ retcode = 0 if not recursive and isDirectory(rootFile,pathSplit): if pathSplit == []: logging.warning(FILE_REMOVE_ERROR.format(rootFile.GetName())) retcode += 1 else: logging.warning(DIRECTORY_REMOVE_ERROR.format(pathSplit[-1])) retcode += 1 else: if interactive: if pathSplit != []: answer = _input(ASK_OBJECT_REMOVE \ .format("/".join(pathSplit),rootFile.GetName())) else: answer = _input(ASK_FILE_REMOVE \ .format(rootFile.GetName())) remove = answer.lower() == 'y' else: remove = True if remove: if pathSplit != []: retcode += deleteObject(rootFile,pathSplit) else: rootFile.Close() os.remove(rootFile.GetName()) return retcode # End of functions shared by rootcp, rootmv and rootrm ########## ########## # Help strings for ROOT command line tools # Arguments SOURCE_HELP = "path of the source." SOURCES_HELP = "path of the source(s)." DEST_HELP = "path of the destination." # Options COMPRESS_HELP = \ """change the compression settings of the destination file (if not already existing).""" INTERACTIVE_HELP = "prompt before every removal." RECREATE_HELP = "recreate the destination file." RECURSIVE_HELP = "recurse inside directories" REPLACE_HELP = "replace object if already existing" # End of help strings ########## ########## # ROOTBROWSE def _openBrowser(rootFile=None): browser = ROOT.TBrowser() _input("Press enter to exit.") def rootBrowse(fileName=None): if fileName: rootFile = openROOTFile(fileName) if not rootFile: return 1 _openBrowser(rootFile) rootFile.Close() else: _openBrowser() return 0 # End of ROOTBROWSE ########## ########## # ROOTCP def _copyObjects(fileName, pathSplitList, destFile, destPathSplit, oneFile, \ recursive, replace): retcode = 0 destFileName = destFile.GetName() rootFile = openROOTFile(fileName) \ if fileName != destFileName else \ destFile if not rootFile: return 1 ROOT.gROOT.GetListOfFiles().Remove(rootFile) # Fast copy necessity for pathSplit in pathSplitList: oneSource = oneFile and len(pathSplitList)==1 retcode += copyRootObject(rootFile, pathSplit, destFile, destPathSplit, \ oneSource, recursive, replace) if fileName != destFileName: rootFile.Close() return retcode def rootCp(sourceList, destFileName, destPathSplit, \ compress=None, recreate=False, recursive=False, replace=False): # Check arguments if sourceList == [] or destFileName == "": return 1 if recreate and destFileName in [n[0] for n in sourceList]: logging.error("cannot recreate destination file if this is also a source file") return 1 # Open destination file destFile = openROOTFileCompress(destFileName, compress, recreate) if not destFile: return 1 ROOT.gROOT.GetListOfFiles().Remove(destFile) # Fast copy necessity # Loop on the root files retcode = 0 for fileName, pathSplitList in sourceList: retcode += _copyObjects(fileName, pathSplitList, destFile, destPathSplit, \ len(sourceList)==1, recursive, replace) destFile.Close() return retcode # End of ROOTCP ########## ########## # ROOTEVENTSELECTOR def _setBranchStatus(tree,branchSelectionString,status=0): """This is used by _copyTreeSubset() to turn on/off branches""" for branchToModify in branchSelectionString.split(","): logging.info("Setting branch status to %d for %s"%(status,branchToModify) ) tree.SetBranchStatus(branchToModify,status) return tree def _copyTreeSubset(sourceFile,sourcePathSplit,destFile,destPathSplit,firstEvent,lastEvent,selectionString, branchinclude, branchexclude): """Copy a subset of the tree from (sourceFile,sourcePathSplit) to (destFile,destPathSplit) according to options in optDict""" retcode = changeDirectory(sourceFile,sourcePathSplit[:-1]) if retcode != 0: return retcode bigTree = getFromDirectory(sourcePathSplit[-1]) nbrEntries = bigTree.GetEntries() # changeDirectory for the small tree not to be memory-resident retcode = changeDirectory(destFile,destPathSplit) if retcode != 0: return retcode if lastEvent == -1: lastEvent = nbrEntries-1 numberOfEntries = (lastEvent-firstEvent)+1 # "Slim" tree by removing branches - # This is done after the skimming to allow for the user to skim on a # branch they no longer need to keep outputTree = bigTree if branchexclude: _setBranchStatus(outputTree,branchexclude,0) if branchinclude: _setBranchStatus(outputTree,branchinclude,1) if branchexclude or branchinclude: outputTree = outputTree.CloneTree() # "Skim" events based on branch values using selectionString # as well as selecting a range of events by index outputTree = outputTree.CopyTree(selectionString,"",numberOfEntries,firstEvent) outputTree.Write() return retcode def _copyTreeSubsets(fileName, pathSplitList, destFile, destPathSplit, first, last, selectionString, branchinclude, branchexclude): retcode = 0 destFileName = destFile.GetName() rootFile = openROOTFile(fileName) \ if fileName != destFileName else \ destFile if not rootFile: return 1 for pathSplit in pathSplitList: if isTree(rootFile,pathSplit): retcode += _copyTreeSubset(rootFile,pathSplit, \ destFile,destPathSplit,first,last,selectionString,branchinclude, branchexclude) if fileName != destFileName: rootFile.Close() return retcode def rootEventselector(sourceList, destFileName, destPathSplit, \ compress=None, recreate=False, first=0, last=-1, selectionString="", branchinclude="", branchexclude=""): # Check arguments if sourceList == [] or destFileName == "": return 1 if recreate and destFileName in sourceList: logging.error("cannot recreate destination file if this is also a source file") return 1 # Open destination file destFile = openROOTFileCompress(destFileName, compress, recreate) if not destFile: return 1 # Loop on the root file retcode = 0 for fileName, pathSplitList in sourceList: retcode += _copyTreeSubsets(fileName, pathSplitList, destFile, destPathSplit, \ first, last, selectionString, branchinclude, branchexclude) destFile.Close() return retcode # End of ROOTEVENTSELECTOR ########## ########## # ROOTLS # Ansi characters ANSI_BOLD = "\x1B[1m" ANSI_BLUE = "\x1B[34m" ANSI_GREEN = "\x1B[32m" ANSI_END = "\x1B[0m" # Needed for column width calculation ANSI_BOLD_LENGTH = len(ANSI_BOLD+ANSI_END) ANSI_BLUE_LENGTH = len(ANSI_BLUE+ANSI_END) ANSI_GREEN_LENGTH = len(ANSI_GREEN+ANSI_END) # Terminal and platform booleans IS_TERMINAL = sys.stdout.isatty() IS_WIN32 = sys.platform == 'win32' def isSpecial(ansiCode,string): """Use ansi code on 'string' if the output is the terminal of a not Windows platform""" if IS_TERMINAL and not IS_WIN32: return ansiCode+string+ANSI_END else: return string def write(string,indent=0,end=""): """Use sys.stdout.write to write the string with an indentation equal to indent and specifying the end character""" sys.stdout.write(" "*indent+string+end) TREE_TEMPLATE = "{0:{nameWidth}}"+"{1:{titleWidth}}{2:{memoryWidth}}" def _recursifTreePrinter(tree,indent): """Print recursively tree informations""" listOfBranches = tree.GetListOfBranches() if len(listOfBranches) > 0: # Width informations maxCharName = max([len(branch.GetName()) \ for branch in listOfBranches]) maxCharTitle = max([len(branch.GetTitle()) \ for branch in listOfBranches]) dic = { \ "nameWidth":maxCharName+2, \ "titleWidth":maxCharTitle+4, \ "memoryWidth":1} for branch in listOfBranches: # Print loop rec = \ [branch.GetName(), \ "\""+branch.GetTitle()+"\"", \ str(branch.GetTotBytes())] write(TREE_TEMPLATE.format(*rec,**dic),indent,end="\n") _recursifTreePrinter(branch,indent+2) write("") def _prepareTime(time): """Get time in the proper shape ex : 174512 for 17h 45m 12s ex : 094023 for 09h 40m 23s""" time = str(time) time = '000000'+time time = time[len(time)-6:] return time MONTH = {1:'Jan',2:'Feb',3:'Mar',4:'Apr',5:'May',6:'Jun', \ 7:'Jul',8:'Aug',9:'Sep',10:'Oct',11:'Nov',12:'Dec'} LONG_TEMPLATE = \ isSpecial(ANSI_BOLD,"{0:{classWidth}}")+"{1:{timeWidth}}" + \ "{2:{nameWidth}}{3:{titleWidth}}" def _printClusters(tree, indent): clusterStart = 0 nTotClusters = 0 clusterIter = tree.GetClusterIterator(0) clusterStart = clusterIter() write(isSpecial(ANSI_BOLD, "Cluster INCLUSIVE ranges:\n"), indent) while clusterStart < tree.GetEntries(): # here we list the inclusive ranges, therefore we have a -1 clustLine = " - # %d: [%d, %d]\n" % ( nTotClusters, clusterStart, clusterIter.GetNextEntry() - 1) write(clustLine, indent) nTotClusters += 1 clusterStart = clusterIter() write(isSpecial(ANSI_BOLD,"The total number of clusters is %d\n" % nTotClusters), indent) def _rootLsPrintLongLs(keyList,indent,treeListing): """Print a list of Tkey in columns pattern : classname, datetime, name and title""" if len(keyList) > 0: # Width informations maxCharClass = max([len(key.GetClassName()) for key in keyList]) maxCharTime = 12 maxCharName = max([len(key.GetName()) for key in keyList]) dic = { \ "classWidth":maxCharClass+2, \ "timeWidth":maxCharTime+2, \ "nameWidth":maxCharName+2, \ "titleWidth":1} for key in keyList: datime = key.GetDatime() time = datime.GetTime() date = datime.GetDate() year = datime.GetYear() time = _prepareTime(time) rec = \ [key.GetClassName(), \ MONTH[int(str(date)[4:6])]+" " +str(date)[6:]+ \ " "+time[:2]+":"+time[2:4]+" "+str(year)+" ", \ key.GetName(), \ "\""+key.GetTitle()+"\""] write(LONG_TEMPLATE.format(*rec,**dic),indent,end="\n") if treeListing and isTreeKey(key): tree = key.ReadObj() _recursifTreePrinter(tree,indent+2) tree = tree.GetTree() _printClusters(tree, indent+2) if treeListing and isTHnSparseKey(key): hs = key.ReadObj() hs.Print('all') ## # The code of the getTerminalSize function can be found here : # https://gist.github.com/jtriley/1108174 # Thanks jtriley !! import os import shlex import struct import platform import subprocess def getTerminalSize(): """ getTerminalSize() - get width and height of console - works on linux,os x,windows,cygwin(windows) originally retrieved from: http://stackoverflow.com/questions/566746/how-to-get-console-window-width-in-python""" current_os = platform.system() tuple_xy = None if current_os == 'Windows': tuple_xy = _get_terminal_size_windows() if tuple_xy is None: tuple_xy = _get_terminal_size_tput() # needed for window's python in cygwin's xterm! if current_os in ['Linux', 'Darwin'] or current_os.startswith('CYGWIN'): tuple_xy = _get_terminal_size_linux() if tuple_xy is None: #print "default" #_get_terminal_size_windows() or _get_terminal_size_tput don't work tuple_xy = (80, 25) # default value return tuple_xy def _get_terminal_size_windows(): try: from ctypes import windll, create_string_buffer # stdin handle is -10 # stdout handle is -11 # stderr handle is -12 h = windll.kernel32.GetStdHandle(-12) csbi = create_string_buffer(22) res = windll.kernel32.GetConsoleScreenBufferInfo(h, csbi) if res: (bufx, bufy, curx, cury, wattr, left, top, right, bottom, maxx, maxy) = struct.unpack("hhhhHhhhhhh", csbi.raw) sizex = right - left + 1 sizey = bottom - top + 1 return sizex, sizey except: pass def _get_terminal_size_tput(): # get terminal width # src: http://stackoverflow.com/questions/263890/how-do-i-find-the-width-height-of-a-terminal-window try: cols = int(subprocess.check_call(shlex.split('tput cols'))) rows = int(subprocess.check_call(shlex.split('tput lines'))) return (cols, rows) except: pass def _get_terminal_size_linux(): def ioctl_GWINSZ(fd): try: import fcntl import termios cr = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, '1234')) return cr except: pass cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2) if not cr: try: fd = os.open(os.ctermid(), os.O_RDONLY) cr = ioctl_GWINSZ(fd) os.close(fd) except: pass if not cr: try: cr = (os.environ['LINES'], os.environ['COLUMNS']) except: return None return int(cr[1]), int(cr[0]) # End of getTerminalSize code ## def _rootLsPrintSimpleLs(keyList,indent,oneColumn): """Print list of strings in columns - blue for directories - green for trees""" # This code is adapted from the pprint_list function here : # http://stackoverflow.com/questions/25026556/output-list-like-ls # Thanks hawkjo !! if len(keyList) == 0: return (term_width, term_height) = getTerminalSize() term_width = term_width - indent min_chars_between = 2 min_element_width = min( len(key.GetName()) for key in keyList ) \ + min_chars_between max_element_width = max( len(key.GetName()) for key in keyList ) \ + min_chars_between if max_element_width >= term_width: ncol,col_widths = 1,[1] else: # Start with max possible number of columns and reduce until it fits ncol = 1 if oneColumn else min( len(keyList), term_width // min_element_width ) while True: col_widths = \ [ max( len(key.GetName()) + min_chars_between \ for j, key in enumerate(keyList) if j % ncol == i ) \ for i in range(ncol) ] if sum( col_widths ) <= term_width: break else: ncol -= 1 for i, key in enumerate(keyList): if i%ncol == 0: write("",indent) # indentation # Don't add spaces after the last element of the line or of the list if (i+1)%ncol != 0 and i != len(keyList)-1: if not IS_TERMINAL: write( \ key.GetName().ljust(col_widths[i%ncol])) elif isDirectoryKey(keyList[i]): write( \ isSpecial(ANSI_BLUE,key.GetName()).ljust( \ col_widths[i%ncol] + ANSI_BLUE_LENGTH)) elif isTreeKey(keyList[i]): write( \ isSpecial(ANSI_GREEN,key.GetName()).ljust( \ col_widths[i%ncol] + ANSI_GREEN_LENGTH)) else: write(key.GetName().ljust(col_widths[i%ncol])) else: # No spaces after the last element of the line or of the list if not IS_TERMINAL: write(key.GetName()) elif isDirectoryKey(keyList[i]): write(isSpecial(ANSI_BLUE, key.GetName())) elif isTreeKey(keyList[i]): write(isSpecial(ANSI_GREEN, key.GetName())) else: write(key.GetName()) write('\n') def _rootLsPrint(keyList, indent, oneColumn, \ longListing, treeListing): """Print informations given by keyList with a rootLs style chosen with the options""" if longListing or treeListing: \ _rootLsPrintLongLs(keyList, indent, treeListing) else: _rootLsPrintSimpleLs(keyList, indent, oneColumn) def _rootLsProcessFile(fileName, pathSplitList, manySources, indent, \ oneColumn, longListing, treeListing): '''rootls main routine for one file looping over paths in the file sorts out directories and key, and loops over all paths, then forwards to (_rootLsPrintLongLs or _rootLsPrintSimpleLs) - split in _rootLsPrint args: oneColumn (bool): longListing (bool): treeListing (bool): indent (int): how many columns the printout should be indented globally manySources (bool): if more than one file is printed fileName (str): the root file name pathSplitList: a list of subdirectories, each being represented as a list itself with a string per level e.g. [['Method_BDT','BDT']] Returns: retcode (int): 0 in case of success, 1 if the file could not be opened ''' retcode = 0 rootFile = openROOTFile(fileName) if not rootFile: return 1 keyList,dirList = keyClassSpliter(rootFile,pathSplitList) # keyList lists the TKey objects from pathSplitList # dirList is 'just the pathSplitList' for what aren't TKeys if manySources: write("{0} :".format(fileName)+"\n") _rootLsPrint(keyList, indent, oneColumn, longListing, treeListing) # Loop on the directories manyPathSplits = len(pathSplitList) > 1 indentDir = 2 if manyPathSplits else 0 for pathSplit in dirList: keyList = getKeyList(rootFile,pathSplit) keyListSort(keyList) if manyPathSplits: write("{0} :".format("/".join(pathSplit)),indent,end="\n") _rootLsPrint(keyList, indent+indentDir, oneColumn, longListing, treeListing) rootFile.Close() return retcode def rootLs(sourceList, oneColumn=False, longListing=False, treeListing=False): '''rootls main routine for an arbitrary number of files args: oneColumn (bool): longListing (bool): treeListing (bool): sourceList: a list of tuples with one list element per file the first tuple entry being the root file, the second a list of subdirectories, each being represented as a list itself with a string per level e.g. rootls tutorial/tmva/TMVA.root:Method_BDT/BDT turns into [('tutorials/tmva/TMVA.root', [['Method_BDT','BDT']])] returns: retcode (int): 0 in case of success ''' # Check arguments if sourceList == []: return 1 # sort sourceList according to filenames tupleListSort(sourceList) # Loop on the ROOT files retcode = 0 manySources = len(sourceList) > 1 indent = 2 if manySources else 0 for fileName, pathSplitList in sourceList: retcode += _rootLsProcessFile(fileName, pathSplitList, manySources, indent, \ oneColumn, longListing, treeListing) return retcode # End of ROOTLS ########## ########## # ROOTMKDIR MKDIR_ERROR = "cannot create directory '{0}'" def _createDirectories(rootFile,pathSplit,parents): """Same behaviour as createDirectory but allows the possibility to build an whole path recursively with the option \"parents\" """ retcode = 0 lenPathSplit = len(pathSplit) if lenPathSplit == 0: pass elif parents: for i in range(lenPathSplit): currentPathSplit = pathSplit[:i+1] if not (isExisting(rootFile,currentPathSplit) \ and isDirectory(rootFile,currentPathSplit)): retcode += createDirectory(rootFile,currentPathSplit) else: doMkdir = True for i in range(lenPathSplit-1): currentPathSplit = pathSplit[:i+1] if not (isExisting(rootFile,currentPathSplit) \ and isDirectory(rootFile,currentPathSplit)): doMkdir = False break if doMkdir: retcode += createDirectory(rootFile,pathSplit) else: logging.warning(MKDIR_ERROR.format("/".join(pathSplit))) retcode += 1 return retcode def _rootMkdirProcessFile(fileName, pathSplitList, parents): retcode = 0 rootFile = openROOTFile(fileName,"update") if not rootFile: return 1 for pathSplit in pathSplitList: retcode+=_createDirectories(rootFile,pathSplit,parents) rootFile.Close() return retcode def rootMkdir(sourceList, parents=False): # Check arguments if sourceList == []: return 1 # Loop on the ROOT files retcode = 0 for fileName, pathSplitList in sourceList: retcode += _rootMkdirProcessFile(fileName, pathSplitList, parents) return retcode # End of ROOTMKDIR ########## ########## # ROOTMV MOVE_ERROR = "error during copy of {0}, it is not removed from {1}" def _moveObjects(fileName, pathSplitList, destFile, destPathSplit, \ oneFile, interactive): retcode = 0 recursive = True replace = True destFileName = destFile.GetName() rootFile = openROOTFile(fileName,"update") \ if fileName != destFileName else \ destFile if not rootFile: return 1 ROOT.gROOT.GetListOfFiles().Remove(rootFile) # Fast copy necessity for pathSplit in pathSplitList: oneSource = oneFile and len(pathSplitList)==1 retcodeTemp = copyRootObject(rootFile,pathSplit, \ destFile,destPathSplit,oneSource,recursive,replace) if not retcodeTemp: retcode += deleteRootObject(rootFile, pathSplit, interactive, recursive) else: logging.warning(MOVE_ERROR.format("/".join(pathSplit),rootFile.GetName())) retcode += retcodeTemp if fileName != destFileName: rootFile.Close() return retcode def rootMv(sourceList, destFileName, destPathSplit, compress=None, \ interactive=False, recreate=False): # Check arguments if sourceList == [] or destFileName == "": return 1 if recreate and destFileName in sourceList: logging.error("cannot recreate destination file if this is also a source file") return 1 # Open destination file destFile = openROOTFileCompress(destFileName,compress,recreate) if not destFile: return 1 ROOT.gROOT.GetListOfFiles().Remove(destFile) # Fast copy necessity # Loop on the root files retcode = 0 for fileName, pathSplitList in sourceList: retcode += _moveObjects(fileName, pathSplitList, destFile, destPathSplit, \ len(sourceList)==1, interactive) destFile.Close() return retcode # End of ROOTMV ########## ########## # ROOTPRINT def _keyListExtended(rootFile,pathSplitList): keyList,dirList = keyClassSpliter(rootFile,pathSplitList) for pathSplit in dirList: keyList.extend(getKeyList(rootFile,pathSplit)) keyList = [key for key in keyList if not isDirectoryKey(key)] keyListSort(keyList) return keyList def rootPrint(sourceList, directoryOption = None, divideOption = None, drawOption = "", formatOption = None, \ outputOption = None, sizeOption = None, styleOption = None, verboseOption = False): # Check arguments if sourceList == []: return 1 tupleListSort(sourceList) # Don't open windows ROOT.gROOT.SetBatch() # (Style option) if styleOption: ROOT.gInterpreter.ProcessLine(".x {0}".format(styleOption)) # (Verbose option) if not verboseOption: ROOT.gErrorIgnoreLevel = 9999 # Initialize the canvas (Size option) if sizeOption: try: width,height = sizeOption.split("x") width = int(width) height = int(height) except ValueError: logging.warning("canvas size is on a wrong format") return 1 canvas = ROOT.TCanvas("canvas","canvas",width,height) else: canvas = ROOT.TCanvas("canvas") # Divide the canvas (Divide option) if divideOption: try: x,y = divideOption.split(",") x = int(x) y = int(y) except ValueError: logging.warning("divide is on a wrong format") return 1 canvas.Divide(x,y) caseNumber = x*y # Take the format of the output file (formatOutput option) if not formatOption and outputOption: fileName = outputOption fileFormat = fileName.split(".")[-1] formatOption = fileFormat # Use pdf as default format if not formatOption: formatOption = "pdf" # Create the output directory (directory option) if directoryOption: if not os.path.isdir(os.path.join(os.getcwd(),directoryOption)): os.mkdir(directoryOption) # Make the output name, begin to print (output option) if outputOption: if formatOption in ['ps','pdf']: outputFileName = outputOption if directoryOption: outputFileName = \ directoryOption + "/" + outputFileName canvas.Print(outputFileName+"[",formatOption) else: logging.warning("can't merge pictures, only postscript or pdf files") return 1 # Loop on the root files retcode = 0 objDrawnNumber = 0 openRootFiles = [] for fileName, pathSplitList in sourceList: rootFile = openROOTFile(fileName) if not rootFile: retcode += 1 continue openRootFiles.append(rootFile) # Fill the key list (almost the same as in root) keyList = _keyListExtended(rootFile,pathSplitList) for key in keyList: if isTreeKey(key): pass else: if divideOption: canvas.cd(objDrawnNumber%caseNumber + 1) objDrawnNumber += 1 obj = key.ReadObj() obj.Draw(drawOption) if divideOption: if objDrawnNumber%caseNumber == 0: if not outputOption: outputFileName = str(objDrawnNumber//caseNumber)+"."+formatOption if directoryOption: outputFileName = os.path.join( \ directoryOption,outputFileName) canvas.Print(outputFileName,formatOption) canvas.Clear() canvas.Divide(x,y) else: if not outputOption: outputFileName = key.GetName() + "." +formatOption if directoryOption: outputFileName = os.path.join( \ directoryOption,outputFileName) if outputOption or formatOption == 'pdf': objTitle = "Title:"+key.GetClassName()+" : "+key.GetTitle() canvas.Print(outputFileName,objTitle) else: canvas.Print(outputFileName,formatOption) # Last page (divideOption) if divideOption: if objDrawnNumber%caseNumber != 0: if not outputOption: outputFileName = str(objDrawnNumber//caseNumber + 1)+"."+formatOption if directoryOption: outputFileName = os.path.join(directoryOption,outputFileName) canvas.Print(outputFileName,formatOption) # End to print (output option) if outputOption: if not divideOption: canvas.Print(outputFileName+"]",objTitle) else: canvas.Print(outputFileName+"]") # Close ROOT files map(lambda rootFile: rootFile.Close(),openRootFiles) return retcode # End of ROOTPRINT ########## ########## # ROOTRM def _removeObjects(fileName, pathSplitList, interactive=False, recursive=False): retcode = 0 rootFile = openROOTFile(fileName,"update") if not rootFile: return 1 for pathSplit in pathSplitList: retcode += deleteRootObject(rootFile, pathSplit, interactive, recursive) rootFile.Close() return retcode def rootRm(sourceList, interactive=False, recursive=False): # Check arguments if sourceList == []: return 1 # Loop on the root files retcode = 0 for fileName, pathSplitList in sourceList: retcode += _removeObjects(fileName, pathSplitList, interactive, recursive) return retcode # End of ROOTRM ##########