eic-smear/eicpy/build.py

0001 #!/usr/bin/env python
0002
0003 """Script for building ROOT trees.
0004
0005 Thomas Burton, BNL, 5th September 2013, tpb@bnl.gov
0006
0007 For usage run
0008     build.py --help
0009
0010 Prerequisites:
0011
0012 -- ROOT installed with Python support (active by default for recent versions)
0013 -- ROOT.py should be accessible via $PYTHONPATH
0014 -- libeicsmear should be accessible vi $LD_LIBRARY_PATH
0015
0016 If in doubt, adding the ROOT library directory (e.g. $ROOTSYS/lib)
0017 to both LD_LIBRARY_PATH and PYTHONPATH should suffice.
0018
0019 """
0020
0021 import os
0022 import Queue
0023
0024 # Names of command to zip and unzip each type of zipped file
0025 ZIP =   {'.gz': 'gzip', '.bz2': 'bzip2'}
0026 UNZIP = {'.gz': 'gunzip', '.bz2': 'bunzip2'}
0027
0028 class File:
0029     """Processes an input file into a ROOT tree file."""
0030
0031     # Use Queue to manage threading.
0032     # This Queue stores all File objects.
0033     queue = Queue.Queue()
0034
0035     # List of allowed input file extensions
0036     supported_extensions = {'.txt', '.dat'}
0037
0038     def __init__(self, filename, outdir, nevents, rezip):
0039         """Constructor.
0040
0041         Initialise with the input file and output information.
0042         Determines whether the input file is zipped.
0043
0044         """
0045         name, extension = os.path.splitext(filename)
0046         # self.zipext stores the zipped extension if the file was zipped
0047         # initially, or None if it was not zipped.
0048         if extension in ZIP:
0049             self.zipext = extension
0050             self.rezip = rezip
0051             # Get the name and extension now that the zip extension is gone
0052             name, extension = os.path.splitext(name)
0053         else:
0054             self.zipext = None
0055             self.rezip = False
0056         self.name = name # Without extension
0057         self.ext = extension # File extension, not zipped extension
0058         self.outdir = outdir
0059         self.nevents = nevents
0060
0061     def process(self):
0062         """Build the tree for the input the File was initialised with.
0063
0064         If the file is zipped, unzip it first and rezip it after making
0065         the ROOT tree (unless requested not to rezip).
0066
0067         """
0068         import subprocess
0069         if self.ext not in File.supported_extensions:
0070             return
0071         fullname = ''.join([self.name, self.ext])
0072         # Unzip a zipped file
0073         if self.zipext:
0074             zipped_name = ''.join([fullname, self.zipext])
0075             # subprocess.call should return 0 if unzipping is successful
0076             unzipped = not subprocess.call(
0077                                 [UNZIP[self.zipext], '-v', zipped_name])
0078         else:
0079             unzipped = False
0080         # Catch any errors from tree building to we always rezip if asked to
0081         try:
0082             # Here's what we actually came to do!
0083             ROOT.BuildTree(fullname, self.outdir, self.nevents)
0084         except:
0085             print 'Error encountered building tree'
0086         if unzipped and self.rezip and self.zipext:
0087             print 'Rezipping', fullname
0088             subprocess.call([ZIP[self.zipext], '-v', fullname])
0089
0090
0091 def processor():
0092     """Function run by each thread.
0093
0094     Pops Files from the queue and processes them until there are no files
0095     remaining.
0096
0097     """
0098     while True:
0099         file = File.queue.get()     # Get the next file from the Queue
0100         file.process()              # Process the file
0101         File.queue.task_done()      # Inform the Queue that file is done
0102
0103 def build_list(filelist, outdir='.', nevents=-1, nthreads = 1, rezip=True):
0104     """Build ROOT tree files from all the listed files.
0105
0106     Arguments:
0107     filelist -- a list of file names
0108     outdir -- the directory to which to write ROOT files [current directory]
0109     nevents -- the maximum number of events per file to process [all]
0110     nthreads -- the maximum number of threads permitted to run [1]
0111
0112     Zipped files (either gzip or bzip2) are unzipped then rezipped
0113     after the ROOT file has been made.
0114
0115     """
0116     import threading
0117     # Generate our threads, each running the processing function
0118     for i in range(nthreads):
0119         t = threading.Thread(target=processor)
0120         t.daemon = True
0121         t.start()
0122     # Populate the Queue with Files
0123     for i in filelist:
0124         File.queue.put(File(i, outdir, nevents, rezip))
0125     # Wait for all the Queue elements to finish processing
0126     File.queue.join()
0127
0128 def parse():
0129     """Process command line arguments and return argparse Namespace."""
0130     from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
0131     parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter,
0132         description='build trees from files in a list or directory')
0133     parser.add_argument('source', help='name of file list or directory')
0134     parser.add_argument('-o', '--outdir', default='.', help='output directory')
0135     parser.add_argument('-e', '--events', type=int, default=-1,
0136                         help='number of events per file')
0137     parser.add_argument('-p', '--processes', type=int, default=1,
0138                         help='number of parallel processes')
0139     parser.add_argument('--norezip', action='store_true',
0140                         help='do not rezip files')
0141     return parser.parse_args()
0142
0143 # Execute buildlist on all the files in a directory, the current
0144 # directory by default
0145 if __name__ == "__main__":
0146     """Executes buildlist on all the files in a list or directory."""
0147     # Parse command line options
0148     options = parse()
0149     # Import ROOT after parsing arguments so build.py --help will still
0150     # work even if ROOT cannot be located.
0151     try:
0152         import ROOT
0153         # Try to tell ROOT to ignore command line arguments.
0154         # Otherwise ROOT will intercept command line arguments and
0155         # build.py --help won't work. This isn't supported in older versions
0156         # of ROOT, so don't complain if it can't be done.
0157         try:
0158             ROOT.PyConfig.IgnoreCommandLineOptions = True
0159         except AttributeError:
0160             pass
0161         # Load libeicsmear, raise and error if it can't be found
0162         if ROOT.gSystem.Load('libeicsmear') < 0:
0163             raise IOError('libeicsmear could not be located')
0164     # If importing ROOT and libeicsmear failed print an error and quit
0165     except Exception as error:
0166         print 'Error:', error
0167         quit()
0168     # Try to get list of file names from an input file
0169     if os.path.isfile(options.source):
0170         with open(options.source) as file:
0171             files = file.read().splitlines()
0172     # Try to get list of file names from all files in a directory
0173     elif os.path.isdir(options.source):
0174         files = [os.path.join(options.source, file)
0175                  for file in os.listdir(options.source)]
0176     # Got some bad input, complain
0177     else:
0178         print options.source, 'not recognized... quitting'
0179         quit()
0180     # Build everything
0181     build_list(files, options.outdir, options.events, options.processes,
0182                not options.norezip)