diff --git a/main/python/cmdLineUtils.py b/main/python/cmdLineUtils.py index f3af283b197db..400afd5a767d4 100644 --- a/main/python/cmdLineUtils.py +++ b/main/python/cmdLineUtils.py @@ -1,1255 +1,7 @@ -#!/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 -from time import sleep -from itertools import zip_longest - -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) - if not cl: - logging.warning("Unknown class to ROOT: " + classname) - return False - 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) - if not cl: - logging.warning("Unknown class to ROOT: " + classname) - return False - 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) - if not cl: - logging.warning("Unknown class to ROOT: " + classname) - return False - 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 keyClassSplitter(rootFile, pathSplitList): - """ - Return a list of directories and a list of keys corresponding - to the other objects, for rootls and rootprint 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, - 'recursiveListing'; 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 = "{0} '{1}' will be copied but not its subdirectories (if any). Use the -r option if you need 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 - if not recursiveOption: - if sourcePathSplit == []: - logging.warning(OMITTING_ERROR.format("file", sourceFile.GetName())) - retcode += 1 - elif isDirectory(sourceFile, sourcePathSplit): - logging.warning(OMITTING_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 : 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 + sourceFile, + sourcePathSplit + [objectName], + destFile, + destPathSplit + [objectName], + replaceOption, ) - 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) - retcodeTemp = deleteObject(destFile, destPathSplit + [setName]) - if retcodeTemp: - retcode += retcodeTemp - continue - else: - if isinstance(obj, ROOT.TNamed): - 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 != "": - if isinstance(obj, ROOT.TNamed): - obj.SetName(setName) - else: - if isinstance(obj, ROOT.TNamed): - 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 -########## - -########## -# 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 -########## - -########## -# 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, recursive = False): - prefixList = [] - keyList, dirList = keyClassSplitter(rootFile, pathSplitList) - for pathSplit in dirList: - keyList.extend(getKeyList(rootFile, pathSplit)) - subList = [key for key in keyList if isDirectoryKey(key)] - keyList = [key for key in keyList if not isDirectoryKey(key)] - prefixList = ["" for key in keyList] - if recursive: - for subdir in subList: - subkeyList, subprefixList = _keyListExtended(ROOT.gDirectory.Get(subdir.GetName()), pathSplitList, recursive) - keyList.extend(subkeyList) - prefixList.extend([subdir.GetName() + "_" + prefix for prefix in subprefixList]) - if recursive: - keyList, prefixList = (list(t) for t in zip(*sorted(zip(keyList, prefixList), key=lambda x: x[0].GetName().lower()))) - else: - keyListSort(keyList) - return keyList, prefixList - - -def rootPrint( - sourceList, - directoryOption=None, - divideOption=None, - drawOption="", - formatOption=None, - outputOption=None, - sizeOption=None, - styleOption=None, - verboseOption=False, - recursiveOption=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, prefixList = _keyListExtended(rootFile, pathSplitList, recursiveOption) - for k, key in enumerate(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: - prefix = prefixList[k] - if not outputOption: - outputFileName = prefix + 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 -##########