diana-rootfilehandler.cxx

Go to the documentation of this file.

#include <TFile.h>
#include "QFileList.hh"
#include <string.h>
#include <stdlib.h>
#include <unistd.h>
#include <stdio.h>
#include <getopt.h>
#include <sys/types.h>
#include <TFile.h>
#include "QAliases.hh"
#include "QEventLabel.hh"
#include <TKey.h>
#include <string>
#include <iostream>
#include <TError.h>
#include <set>
#include <list>
#include <map>
#include <unistd.h>
#include "QTree.hh"
#ifndef __USING_ROOT6__
  #include <TCint.h>
#endif
#include <TFileMerger.h>
#include <TInterpreter.h>
#include "QFrameWorkConfig.hh"
#include "QGlobalDataManager.hh"
#include "QStringHandler.hh"
 
#include <TROOT.h>
 
/* customized structure for command line parameters */
struct arg_t
{
  bool l;
  bool b;
  bool c;
  bool t;
  char * r;
  bool h;
  bool v;
  char *C;
  char *S;
  char *a;
  int optind;
};
 
/* function prototypes */
struct arg_t * Cmdline(int, char **);
void usage(char *);
void free_args(struct arg_t *);
struct arg_t *Cmdline(int argc, char *argv[]);
std::string RemoveDuplicates(const std::string& input,char d);
std::vector<std::string> GetTrees(TFile* file);
 
//int gErrorIgnoreLevel = kError;
 
void ParseDirAndFileName(const std::string& iinput, std::string& dir, std::string& file)
{
    std::string input = RemoveDuplicates(iinput,'/');
    size_t slash = input.find_last_of("/");
    if(slash != std::string::npos) {
        dir = input.substr(0,slash+1);
        file = input.substr(slash+1,std::string::npos);
    } else {
       char cwd[1024];
       getcwd(cwd, sizeof(cwd));
       dir = cwd;
       dir += "/";
       file = input;
    }
}
 
std::string RemoveDuplicates(const std::string& input,char d)
{
    std::string output;
    if(input.empty()) return output;
    char prev = input[0];
    output.push_back(prev);
    for(size_t i = 1; i < input.size(); i++) {
        char curr = input[i];
        if(!(curr == prev && curr == d)) {
            output.push_back(curr);
        }
        prev = curr;
    }
    return output;
}
 
int main(int argc, char* argv[])
{
    arg_t* args = Cmdline(argc,argv);
 
    int optind = args->optind;
    if(optind == 0) {
        usage(argv[0]);
    }
 
    #ifdef __USING_ROOT6__
    // MV FIXME
//        ::ROOT::GetROOT()->GetInterpreter()->EnableAutoLoading();
    #else
        TCint* cint = new TCint("pino","pino");
        cint->EnableAutoLoading();
    #endif
 
    // gather filenames
    std::set<std::string> filenames;
    for(int f = args->optind; f < argc; f++) {
        std::string fileExtension = "";
        std::string inFile = argv[f];
        size_t dotpos =  inFile.find_last_of(".");
        if(dotpos != std::string::npos)
            fileExtension = inFile.substr( dotpos+1 );
        if(fileExtension == "list") {
            std::list<std::string> list = QFileList::Read(inFile,"");
            if(list.empty()) {
                std::cerr<<"Error: problems reading file list "<<inFile<<std::endl;
                exit(2);
            }
            std::list<std::string>::const_iterator iter = list.begin();
            while(iter != list.end()) {
                filenames.insert(*iter);
                iter++;
            }
        } else {
            filenames.insert(inFile);
        }
    }
    // check they are all good files and group by dir
    std::set<std::string>::const_iterator fiter = filenames.begin();
    std::map<std::string, std::set<std::string> > dirfiles;
    while(fiter != filenames.end()) {
        TFile tfile(fiter->c_str());
        if(tfile.IsZombie()) {  
            std::cerr<<"Error: file "<<fiter->c_str()<<" cannot be opened"<<std::endl;
            exit(2);
        }
        std::string dir,file;
        ParseDirAndFileName(*fiter,dir,file);
        dirfiles[dir].insert(file);
        fiter++;
    }
 
 
 
    std::map<std::string, std::set<std::string> >::const_iterator dfiter = dirfiles.begin();
    while(dfiter != dirfiles.end()) { 
        filenames = dfiter->second;
        fiter = filenames.begin();
        std::cout<<"== Processing directory: "<<dfiter->first<<std::endl;
        while(fiter != filenames.end()) {
            // list content of each file
            std::string filename = dfiter->first +  *fiter;
            std::cout<<*fiter<<": ";
            std::cout.flush();
            bool found = false;
            if(args->v || args->t) {
                TFile file(filename.c_str());
                std::vector<std::string> trees = GetTrees(&file);
                std::cout<<" QTrees { ";
                for(size_t t = 0; t < trees.size(); t++) {
                    QTree* qtree = (QTree*)file.Get(trees[t].c_str());
                    std::cout<<trees[t]<<"[";
                    if(args->v) std::cout<<qtree->GetSoftwareRevision()<<"] ";
                    else if(args->t) std::cout<<qtree->GetVersionTag()<<"] ";
                }
                std::cout<<"}";
            }
            else if(args->l) {
                TFile file(filename.c_str());
                std::vector<std::string> trees = GetTrees(&file);
                if(args->b) {
                    std::cout<<" TBranches { ";
                    for(size_t t = 0; t < trees.size(); t++) {
                        QTree* qtree = (QTree*)file.Get(trees[t].c_str());
                        TObjArray* leaflist = qtree->GetListOfBranches();
                        TIter next(leaflist);
                        while(TObject* leaf = next()) {
                            std::cout<<leaf->GetName()<<" ";
                        }
                    }
                    std::cout<<"}";
                } else {
                    std::cout<<" QTrees { ";
                    for(size_t t = 0; t < trees.size(); t++) {
                        std::cout<<trees[t]<<" ";
                    }
                    std::cout<<"}";
                }
            } else if(args->a) {
               Diana::QAliases aliases; 
               QError err = aliases.FillFromFile(args->a);
               if(err != QERR_SUCCESS) {
                   std::stringstream str;
                   str<<err;
                   std::cout<<str.str().c_str()<<std::endl;;
                   return 1;
               }
               TFile file(filename.c_str(),"UPDATE");
               std::vector<std::string> trees = GetTrees(&file);
               for(size_t t = 0; t < trees.size(); t++) {
                   QTree* qtree = (QTree*)file.Get(trees[t].c_str());
                   TObjArray* leaflist = qtree->GetListOfBranches();
                   TIter next(leaflist);
                   while(TObject* leaf = next()) {
                       std::string labelstr = leaf->GetName();
                       labelstr = labelstr.substr(0,labelstr.size()-1);
                       Diana::QEventLabel label(labelstr);
                       std::map<std::string, std::string> aliasMap = aliases.Find(label);
                       std::map<std::string, std::string>::const_iterator aiter = aliasMap.begin();
 
 
                       while(aiter != aliasMap.end()) {
                           std::string path;
                           if(aiter->second =="") {
                               path = labelstr;
                           } else {
                               path = labelstr+std::string(".")+aiter->second;
 
                           }
                           const char* al = aiter->first.c_str();
                           const char* ap = path.c_str();
                       //    printf("Adding alias %s for %s\n",al,ap);
                           qtree->SetAlias(al,ap);
                           aiter++;
                       }
                   }
               }
               file.Write("",TObject::kOverwrite);
 
            }else if(args->C || args->S) {
                Diana::QGlobalDataManager dm;
                Diana::GlobalHandle<QFrameWorkConfig> fwcHandle("Config");
                dm.Get("Diana",&fwcHandle,filename);
                if(fwcHandle.IsValid()) {
                    const QFrameWorkConfig& fwc = fwcHandle.Get();
                    if(args->S) {
                        std::string configfilename = args->S;
                        configfilename += "/";
                        configfilename += *fiter;
                        configfilename += ".cfg";
                        bool ok = fwc.SaveAllSequences(configfilename);
                        if(!ok) exit(3);
                        std::cout<<"Created \""<<configfilename<<"\"";
                    } else if(args->C) {
                        std::string directory = args->C;
                        std::string prefix = *fiter;
                        prefix += "_";
                        bool  ok = fwc.SaveAllFiles(directory,prefix);
                        if(!ok) exit(3);
                        std::cout<<"Created config files in dir \""<<directory<<"\"";
                    }
                }
 
            } else if(args->r) {
                std::list<std::string> delObjects;
                QStringHandler::Split(args->r,delObjects,',');
                TFile file(filename.c_str(),"UPDATE");
                if(args->b) {
                    // remove TBranch
                    std::vector<std::string> trees = GetTrees(&file);
                    std::vector<QTree*> qtreesToUpdate;
                    for(std::list<std::string>::const_iterator d = delObjects.begin(); d != delObjects.end(); d++) {
                        TString delObject = d->c_str();
                        bool foundThis = false;
                        for(size_t t = 0; t < trees.size(); t++) {
                            QTree* qtree = (QTree*)file.Get(trees[t].c_str());
                            if(qtree->FindBranch(delObject)) {
                                found = true;
                                foundThis = true;
                                std::string bname = delObject.Data();
                                qtree->SetBranchStatus(bname.c_str(),0);
                                bname += "*";
                                qtree->SetBranchStatus(bname.c_str(),0);
                                if(std::find(qtreesToUpdate.begin(),qtreesToUpdate.end(),qtree) == qtreesToUpdate.end()) {
                                    qtreesToUpdate.push_back(qtree);
                                }
                            }                         
                        }
                        if(!foundThis)  {
                            std::cout<<"TBranch "<<delObject<<" to be removed not found ";
                            std::cout.flush();
                        }
 
                    }
                    if(found) {
                        std::string tmpOutput= filename+".d-rfh";
                        TFile outputFile(tmpOutput.c_str(),"RECREATE");
                        for(size_t t = 0; t < qtreesToUpdate.size(); t++) {
                            QTree* qtree = qtreesToUpdate[t];
                            outputFile.cd();
                            qtree->CloneTree();
                            TString delObject = qtree->GetName();
                            file.Delete(delObject+";*");
                        }
                        outputFile.Write();
                        file.Write("",TObject::kOverwrite);
 
                        TFileMerger merg(0);
 
                        merg.AddFile(&file);
                        merg.AddFile(&outputFile);
                        std::string tmpOutput2= filename+".d-rfh2";
                        merg.OutputFile(tmpOutput2.c_str(),1);
                        if(merg.Merge()) {
                              rename(tmpOutput2.c_str(),filename.c_str());
                              remove(tmpOutput.c_str());
                        } else {
                            std::cout<<"Error merging "<<filename<<": its content is unfortunately now split in: "<<filename<<" and "<<tmpOutput<<std::endl;
                            remove(tmpOutput2.c_str());
                        }
 
                    }
                } else {
                    // remove QTree
                    for(std::list<std::string>::const_iterator d = delObjects.begin(); d != delObjects.end(); d++) {
                        TString delObject = d->c_str();
                        if(file.Get(delObject)) {
                            found = true;
                            file.Delete(delObject+";*");
                            TObject* fele = 0;
                            if(QTree* itree = (QTree*) file.Get(QTREE_NAME)) {
                                TList* flist = itree->GetListOfFriends();
                                if(flist) fele = flist->FindObject(delObject.Data());
                                if(fele) itree->GetListOfFriends()->Remove(fele);
                            }
                        } else {
                            std::cout<<"QTree "<<delObject<<" to be removed not found ";
                            std::cout.flush();
                        }
                    }
                    if(found)  {
                        file.Write("",TObject::kOverwrite);
                    }
                }
                std::cout<<std::endl;
            }
 
            if(args->c && !args->b) {
                TFileMerger merg(0);
                TFile file(filename.c_str());
 
                merg.AddFile(&file);
                std::string tmpOutput= filename+".d-rfh";
                merg.OutputFile(tmpOutput.c_str(),1);
                if(merg.Merge()) {
                    rename(tmpOutput.c_str(),filename.c_str());
                } else {
                    std::cout<<"Error merging "<<filename<<std::endl;
                    remove(tmpOutput.c_str());
                }
            }
            fiter++;
            std::cout<<std::endl;
        } // end loop on files
        dfiter++;
    } // end loop on dirs
 
    return 0;
}
 
 
 
 
 
void usage(char *executable)
{
    printf("usage: %s [ -alcrSCh ] .root and/or .list filenames \n", executable);
    printf("  [ -l ] ");
    printf("[ --list  ] ");
    printf("\n");
    printf("         list QTrees (or their TBranches) in input files\n");
 
 
    printf("  [ -t ] ");
    printf("[ --tag  ] ");
    printf("\n");
    printf("         list version tag of QTrees in input files\n");
 
    printf("  [ -v ] ");
    printf("[ --version  ] ");
    printf("\n");
    printf("         list software revision of QTrees in input files\n");
 
    printf("  [ -b ] ");
    printf("[ --branch  ] ");
    printf("\n");
    printf("         let -l and -r options act on TBranches instead of QTrees\n");
 
    printf("  [ -c ] ");
    printf("[ --compact  ] ");
    printf("\n");
    printf("         optimize file size (takes time)\n");
 
    printf("  [ -r ] ");
    printf("[ --remove  ] object_name1,object_name2,... ");
    printf("\n");
    printf("         remove QTree (or TBranch) with specified name\n");
 
    printf("  [ -C ] ");
    printf("[ --configfile  ] outputdir ");
    printf("\n");
    printf("         extract config files containing all sequences used to produce each file\n");
 
    printf("  [ -S ] ");
    printf("[ --single-configfile  ] outputdir ");
    printf("\n");
    printf("         extract a single config file containing all sequences used to produce each file\n");
 
    printf("  [ -a ] ");
    printf("[ --aliasfile  ] aliasfile ");
    printf("\n");
    printf("         apply aliases from aliasfile\n");
 
    printf("  [ -h ] ");
    printf("[ --help  ] ");
    printf("\n");
    printf("         Display help information.\n");
 
    printf("\nexample 1: (lists QTrees in root files and lists files in current dir)\n");
    printf("\n-> diana-rootfilehandler -l *.root *.list\n\n");
 
    printf("example 2: (get all executed sequences in a single config file per root file)\n");
    printf("\n-> diana-rootfilehandler -S cfg  *.root *.list\n\n");
 
    exit(1);
}
 
/*----------------------------------------------------------------------------
 **
 ** free_args()
 **
 ** Call this to free the memory that was dynamically allocated by the parser.
 **
 **--------------------------------------------------------------------------*/
 
void free_args(struct arg_t *my_args)
{
    if (my_args->r != NULL) free(my_args->r);
    if (my_args->C != NULL) free(my_args->C);
    if (my_args->S != NULL) free(my_args->S);
    if (my_args->a != NULL) free(my_args->a);
    free(my_args);
}
 
/*----------------------------------------------------------------------------
 **
 ** Cmdline()
 **
 ** Parse the argv array into the command line structure
 **
 **--------------------------------------------------------------------------*/
 
struct arg_t *Cmdline(int argc, char *argv[])
{
    extern char *optarg;
    extern int optind;
    int option_index = 0;
    int c;
    struct arg_t *my_args;
    int errflg = 0;
 
    static struct option long_options[] =
    {
        {"list", 0, 0, 'l'},
        {"version", 0, 0, 'v'},
        {"tag", 0, 0, 't'},
        {"branch", 0, 0, 'b'},
        {"compact", 0, 0, 'c'},
        {"remove", 1, 0, 'r'},
        {"help", 0, 0, 'h'},
        {"configfile", 1, 0, 'C'},
        {"single-configfile", 1, 0, 'S'},
        {0, 0, 0, 0}
    };
 
    my_args = (struct arg_t *) malloc (sizeof(struct arg_t));
 
    my_args->l = false;
    my_args->b = false;
    my_args->r = NULL;
    my_args->c = false;
    my_args->h = false;
    my_args->v = false;
    my_args->t = false;
    my_args->C = NULL;
    my_args->S = NULL;
    my_args->a = NULL;
 
    while ((c = getopt_long(argc, argv, "a:lvtbcr:hC:S:", long_options, &option_index)) != EOF)
    {
        switch(c)
        {
            case 'l': 
                my_args->l = true;
                break;
            case 'v': 
                my_args->v = true;
                break;
            case 't': 
                my_args->t = true;
                break;
            case 'b': 
                my_args->b = true;
                break;
            case 'c': 
                my_args->c = true;
                break;
            case 'r': 
                my_args->r = strdup(optarg);
                break;
            case 'h': 
                my_args->h = true;
                usage(argv[0]);
                break;
            case 'C': 
                my_args->C = strdup(optarg);
                break;
            case 'S': 
                my_args->S = strdup(optarg);
                break;
            case 'a': 
                my_args->a = strdup(optarg);
                break;
            default:
                usage(argv[0]);
 
        }
    } /* while */
 
    if (errflg)
        usage(argv[0]);
 
    if (optind >= argc)
        my_args->optind = 0;
    else
        my_args->optind = optind;
    return my_args;
}
 
std::vector<std::string> GetTrees(TFile* file)
{
    TList* keys = file->GetListOfKeys();
    std::vector<std::string> trees;
    if(keys) {
        TIter next(keys);
        while(TKey* key = (TKey*) next()) {
            std::string className = key->GetClassName();
            if(className != "QTree") continue;
            std::string objectName = key->GetName();
            if(objectName.find("qtree") == 0) {
                trees.push_back(key->GetName());
            }
        }
    }
    return trees;
}