catGetDatasetInfo

#!/usr/bin/env python

import sys, os

gidMap = {
    'v8-0-0':'1203220026',
    'v7-6-6':'183143795',
    'v7-6-5':'773772766',
    'v7-3-4':'857222646',
    'v7-4-1':'212761185',
    'v7-4-2':'398123591',
    'v7-4-3':'512197961',
    'v7-4-4':'1038794205',
    'v7-4-5':'519625370',
    'v7-4-6':'1087889200',
    'v7-6-1':'1617837617',
    'v7-6-2':'1966334211',
    'v7-6-3':'266478682',
    'v7-6-4':'600658293',
}
    
era = ''    
if len(sys.argv) < 2:
    era = sorted(gidMap)[-1]
    gid = gidMap[era]
else:
    era = sys.argv[1]
    gid = gidMap[sys.argv[1]]
    
print "Using =",era

from ROOT import *

from urllib import urlopen
import csv
gdocbase = "https://docs.google.com/spreadsheets/d/1rWM3AlFKO8IJVaeoQkWZYWwSvicQ1QCXYSzH74QyZqE"
print "Retrieving dataset info from google doc..."
print "Source URL =", gdocbase
csv = list(csv.reader(urlopen(gdocbase + "/pub?gid=%s&single=true&output=csv" % gid).readlines()))

ds = []
nameIdx = csv[0].index("Name")
titleIdx = csv[0].index("Title")
typeIdx = csv[0].index("Type")
xsecIdx = csv[0].index("Cross section (pb)")
nevtIdx = csv[0].index("NEvent")
avgWgtIdx = csv[0].index("Avg GenWgt")
lumiIdx = csv[0].index("luminosity (fb-1)")
colourIdx = csv[0].index("colour")
pathIdx = csv[0].index("Path")
DataSetNameIdx = csv[0].index("DataSetName")
#GlobalTagIdx = csv[0].index("GlobalTag")
#LumiMaskIdx = csv[0].index("LumiMask")
for l in csv[1:]:
    if len(l) == 0 or len(l[0]) == 0: continue
    xsec, nevt, lumi, colour = l[xsecIdx], l[nevtIdx], l[lumiIdx], l[colourIdx]
    avgWgt = l[avgWgtIdx]
    if xsec == "": xsec = "0"
    if nevt == "": nevt = "0"
    if lumi == "": lumi = "0"
    if avgWgt == "": avgWgt = "0"
    ds.append({'title':l[titleIdx],
               'name':l[nameIdx],
               'xsec':float(xsec),
               'nevt':int(nevt), 'avgWgt':float(avgWgt),
               'lumi':float(lumi),
               'colour':eval(colour),
               'path':l[pathIdx],
               'DataSetName':l[DataSetNameIdx],
               'type':l[typeIdx],
               #'GlobalTag':l[GlobalTagIdx],
               #'LumiMask':l[LumiMaskIdx],
               })

outDir = os.environ["CMSSW_BASE"]+"/src/CATTools/CatAnalyzer/data"
if not os.path.exists(outDir+"/dataset"): os.makedirs(outDir+"/dataset")

import json
f = open("%s/dataset/dataset.json" % outDir, "w")
print "Saving dataset info to %s/dataset/dataset.json" % outDir
print>>f, json.dumps(ds, indent=4, sort_keys=True)
f.close()

print "Trying to verify file list..."
from CATTools.CatAnalyzer.libxrd import *
cmd, xrdbase = guessxrd()

## A hack to add prefix for the sites
prefix = ""
hostname = os.environ["HOSTNAME"]
if "sdfarm" in hostname:
    prefix = "root://cms-xrdr.sdfarm.kr:1094//"+xrdbase
##

# since all files are at kisti
if not "sdfarm" in hostname:    
    sys.exit()

storageprefix = "/store/group/CAT"
unmatching = []
emptydirectory = []
print "Listing dataset location"
print "-"*80
for d in ds:
    dataSetName = d['DataSetName']
    dataSetName = dataSetName.split('/')
    if len(dataSetName) < 4: continue

    pdName, sdName = dataSetName[1], dataSetName[2]
    path = "%s/%s/%s_%s" % (storageprefix, pdName, era, sdName)
    pathname = ''
    ## #pathnames = map( lambda x: x, listxrd(path)[0] )
    ## pathnames = listxrd(path)[0]
    ## sorted(pathnames)
    ## print pathnames
    tempVersion = 0
    for ll in listxrd(path)[0]:
        currentVersion = int(ll[-13:-7] + ll[-6:])
        if currentVersion > tempVersion:
            tempVersion = currentVersion
            pathname = ll 
    print pathname

    pp = d['path']
    d['files'] = []
    d['size'] = 0
    for dd in listxrd(pp)[0]:
        files, size = listxrd(dd)
        d['size'] += size
        for ff in files:
            if not ff.endswith(".root"): continue
            ff = ff[len(pp)+1:]
            d['files'].append(ff)
    d['path'] = prefix + pp
    d['files'].sort()
    
    f = open("%s/dataset/dataset_%s.txt" % (outDir, d['name']), "w")
    for key in d:
        if key == "files": continue
        print>>f, "#", key, "=", d[key]
    for ff in d['files']:
        print>>f, os.path.join(d['path'], ff)
    f.close()
    if not pathname in pp:
        unmatching.append(pathname)
    if len(d['files']) == 0:
        emptydirectory.append(pdName)
print "-"*80

if len(unmatching) > 0:
    print 'List of datasets not matching with database'
    for i in unmatching:
        print i
    print "-"*80

if len(emptydirectory) > 0:
    print 'List of datasets with no data'
    for i in emptydirectory:
        print i
    print "-"*80
   
print "Saved dataset location info to %s/dataset/dataset_SampleName.txt" % (outDir)