BeamSpotProducer/scripts/splitter.py

0001 #!/usr/bin/env python3
0002 from __future__ import print_function
0003 import sys,os,subprocess
0004 from CommonMethods import *
0005
0006 class FileObj:
0007     def __init__(self):
0008         self.run       = 0
0009         self.size      = 0
0010         self.fileNames = []
0011
0012 def getRunNumberFromFileName(fileName):
0013     regExp = re.search('(\D+)_(\d+)_(\d+)_(\d+)',fileName)
0014     if not regExp:
0015         return -1
0016     return long(regExp.group(3))
0017
0018
0019
0020 def main():
0021     if len(sys.argv) < 2:
0022         error = "Usage: splitter fromDir"
0023         exit(error)
0024     sourceDir = sys.argv[1] + '/'
0025
0026     fileList = ls(sourceDir,".txt")
0027
0028     fileObjList = {}
0029
0030     totalSize = 0
0031     for fileName in fileList:
0032         runNumber = getRunNumberFromFileName(fileName)
0033         if runNumber not in fileObjList:
0034             fileObjList[runNumber] = FileObj()
0035             fileObjList[runNumber].run = runNumber
0036         fileObjList[runNumber].fileNames.append(fileName)
0037         aCommand  = 'ls -l '+ sourceDir + fileName
0038         output = subprocess.getstatusoutput( aCommand )
0039         fileObjList[runNumber].size += int(output[1].split(' ')[4])
0040         totalSize += int(output[1].split(' ')[4])
0041
0042     sortedKeys = sorted(fileObjList.keys())
0043
0044     split=13
0045
0046     dirSize = 0
0047     tmpList = []
0048     for run in sortedKeys:
0049         dirSize += fileObjList[run].size
0050         tmpList.append(fileObjList[run])
0051         if dirSize > totalSize/split or run == sortedKeys[len(sortedKeys)-1]:
0052             newDir = sourceDir + "Run" + str(tmpList[0].run) + "_" + str(tmpList[len(tmpList)-1].run) + "/"
0053             aCommand  = 'mkdir '+ newDir
0054             output = subprocess.getstatusoutput( aCommand )
0055             print(str(100.*dirSize/totalSize) + "% " + "Run" + str(tmpList[0].run) + "_" + str(tmpList[len(tmpList)-1].run))
0056             for runs in tmpList:
0057                 #print 'cp '+ sourceDir + runs.fileNames[0] + " " + newDir
0058                 cp(sourceDir,newDir,runs.fileNames)
0059             tmpList = []
0060             dirSize = 0
0061
0062
0063
0064
0065     print(totalSize)
0066     print(sortedKeys)
0067     exit("ok")
0068
0069
0070
0071
0072
0073
0074     if not os.path.isdir(destDir):
0075         error = "WARNING: destination directory doesn't exist! Creating it..."
0076         print(error)
0077         os.mkdir(destDir)
0078     copiedFiles = cp(sourceDir,destDir,fileList)
0079
0080     if len(copiedFiles) != len(fileList):
0081         error = "ERROR: I couldn't copy all files from castor"
0082         exit(error)
0083
0084     for fileName in fileList:
0085         fullFileName = destDir + fileName
0086         runNumber = -1;
0087         with open(fullFileName,'r') as file:
0088             for line in file:
0089                 if line.find("Runnumber") != -1:
0090                     tmpRun = int(line.split(' ')[1])
0091                     if runNumber != -1 and tmpRun != runNumber:
0092                         error = "This file (" + fileName + ") contains more than 1 run number! I don't know how to deal with it!"
0093                         exit(error)
0094                     runNumber = int(line.split(' ')[1])
0095         file.close()
0096         newFileName = fileName.replace("None",str(runNumber))
0097         if fileName != newFileName:
0098             aCmd = "mv " + destDir + fileName + " " + destDir + newFileName
0099             print(aCmd)
0100             output =  subprocess.getstatusoutput(aCmd)
0101             if output[0] != 0:
0102                 print(output[1])
0103         else:
0104             print("WARNING couldn't find keyword None in file " + fileName)
0105
0106
0107
0108
0109
0110 if __name__ == "__main__":
0111     main()