HiggsAnalysis/HiggsAnalysis-00-01-00/run/data/datatools.py

'''
This module stores and accesses the datasets
to be used in NtupleWriter.
'''
__version__ = '$Revision: 1$'
__author__  = 'Steffen.Kaiser'

# >>> import datatools
# print docstring: >>> print datatools.__doc__
# print help:      >>> print help(datatools)

import os, sys

class DataLib(object):
    '''Class which stores and finds the datasets'''

    def __init__(self):
        self.itemsitems = []

    def __str__(self):
        return_string = ''
        for i in range(len(self.itemsitems)):
            return_string += str(self.itemsitems[i]) + '\n'
            return_string = return_string.replace('DataSet:','DataSet[%i]:' % i)
        return return_string

    def __len__(self):
        return len(self.itemsitems)

    def check_existence(self, item):
        '''Check if dataset already exists in the library;
        return True if it exists'''
        for dataset in self.itemsitems:
            if(item.compare(dataset) == True):
                return True
        return False
    
    def add(self, item):
        '''Add dataset to the library'''
        assert (self.check_existencecheck_existencecheck_existencecheck_existence(item) == False),\
               str(item) + '\nalready exists'
        self.itemsitems.append(item)

    def add_dataset(self, d_short='', d_long='',path=''):
        '''Create dataset and add it to the library'''
        self.addaddaddadd( DataSet(d_short, d_long, path) )
        
    def get(self, i):
        '''Return dataset i'''
        return self.itemsitems[i]

    def find_dataset(self, description):
        '''Return dataset with the short or long description: description'''
        print 'find dataset with short/long description: ', description

        ds = DataSet()
        for dataset in self.itemsitems:

            if (dataset.get_short() == description) or\
               (dataset.get_long()  == description):
                
                ds = dataset
                break

        if(ds.is_empty() == True):
            print 'Error: Couldn\'t find datset!'
            sys.exit()
        else:
            return ds


class DataSet(object):
    '''Class which stores the information for one dataset'''

    def __init__(self, d_short='', d_long='',path=''):
        self.resetresetresetreset(d_short, d_long, path)

    def __str__(self):
        return_string = 'DataSet:\n'
        return_string += '  d_short: ' + self.d_shortd_short + ',\n'
        return_string += '  d_long:  ' + self.d_longd_long  + ',\n'
        return_string += '  path:    ' + self.pathpath
        return return_string

    def reset(self, d_short, d_long, path):
        assert (type(d_short) is str), 'd_short is not a string'
        assert (type(d_long)  is str), 'd_long is not a string'
        assert (type(path)    is str), 'path is not a string'
        self.d_shortd_short = d_short
        self.d_longd_long  = d_long
        self.pathpath  = path

    def is_empty(self):
        '''Return True if dataset is empty'''
        if(self.d_shortd_short != ''):
            return False
        if(self.d_longd_long != ''):
            return False
        if(self.pathpath != ''):
            return False
        return True

    def compare(self, item):
        '''Compare this dataset to another; return True if d_short or d_long agree'''
        if (self.d_shortd_short == item.d_short) or\
           (self.d_shortd_short == item.d_long):
            return True
        if (self.d_longd_long == item.d_short) or\
           (self.d_longd_long == item.d_long):
            return True
        return False

    def get_short(self):
        '''Return short description of the dataset'''
        return self.d_shortd_short

    def get_long(self):
        '''Return long description of the dataset'''
        return self.d_longd_long

    def get_path(self):
        '''Return path of the dataset'''
        return self.pathpath

    #def get(self):
    #    '''Return dataset'''
    #    return (self.d_short, self.d_long, self.path)



def include_files(path, EventSelector):
    '''Method which appends all .root files in path/ to InputCollections'''

    assert (os.path.exists(path)), path + ' doesn\'t exist'

    #print path
    files=os.listdir(path)

    for filename in files:

        # need to check which is better 'endswith' or 'find'
        #if filename.endswith(".root"):
        if filename.find(".root") != -1:
            #data_string += path + '/' + filename + '\n'
            EventSelector.InputCollections += [path + '/' + filename]
            
            
    print '\n************************************************************\n'
    print ' Runing on the following files: \n'
    
    print 'path:  ', path, '\n'
    print 'files: ', files, '\n'

    print 'InputColl:', EventSelector.InputCollections
    print '\n************************************************************\n'

---