QRootFileWriterBase.cc

Go to the documentation of this file.

#include "QRootFileWriterBase.hh"
#include <iostream>
#include <sstream>
#include <TFile.h>
#include <unistd.h>
#include "QEvent.hh"
#include "QTree.hh"
#include "QAliases.hh"
#include "QGlobalRWFactory.hh"
#include "QEventInspector.hh"
#include "QString.hh"
#include "QDianaR.hh"
#include "QOptions.hh"
#include "QBool.hh"
#include "TKey.h"
 
#include <sys/stat.h>
 
using namespace Diana;
 
//_____________________________________________________________________________
QRootFileWriterBase::QRootFileWriterBase() : QFileWriter(".root") {
  fRootOutput = 0;
  fTree = 0;
}
//_____________________________________________________________________________
void QRootFileWriterBase::Open(const std::string &filename,
                               const QEventInspector &evi) {        
  if(filename == fRWCommon->fReaderCurrentFileName) 
    Panic("Output filename is equal to input filename");
 
  
  fReaderWriterFiles[fRWCommon->fReaderCurrentFileName] = filename;
  // get configuration
  std::string treeDescription = GetString("Description", "Diana tree",false);
  
  fRootOutput = new TFile(filename.c_str(),"RECREATE");  
  if(!fRootOutput || fRootOutput->IsZombie())
    Panic("cannot open file: %s",filename.c_str());
  
  // set the file permissions
  int FilePermissions = strtol(GetString("FilePermissions","0644",false).c_str(),NULL,8);
  chmod(filename.c_str(),FilePermissions);
 
  // convert address to strine
  fRootOutputAddress.fP = fRootOutput;
  fTree = new QTree(treeDescription.c_str());
  
  QTreeInfo& treeInfo = fTree->GetTreeInfo();
  treeInfo.fSoftwareRevision =
    QOptions::GetInstance().GetParameter("General.DIANA_SWREVISION").
    GetString();
  treeInfo.fVersionTag =
    QOptions::GetInstance().GetParameter("General.DIANA_VERSIONTAG").GetString();
  
  const QVdt::QVdt_vector& vec =
    QOptions::GetInstance().GetParameter("General.DIANA_EXTSWREVISION").
    GetVector();
  for(size_t s = 0; s < vec.size(); s++) {
    const QVdt& p = vec[s];
    treeInfo.fExtSWVersion[p.GetName()]=p.GetString();
  }
  
  QAliases aliases = evi.GetAliases();
  std::string diana_install = "";
  try {
    diana_install =  QOptions::GetInstance().GetString("CFGVar.DIANA_INSTALL");
    diana_install += "/";
  } catch(...) {};
  std::string AliasFileName =
    GetString("AliasFileName",diana_install+"cfg/aliases.txt");
  std::string CustomAliasFileName = GetString("CustomAliasFileName","",false);
  if(AliasFileName != "") {
    Debug("Adding aliases from %s",AliasFileName.c_str());
    QError err = aliases.FillFromFile(AliasFileName);
    if(err != QERR_SUCCESS) {
      std::stringstream str;
      str<<err;
      Error("%s",str.str().c_str());
    }
  }
  if(CustomAliasFileName != "") {
    Debug("Adding aliases from %s",CustomAliasFileName.c_str());
    QError err = aliases.FillFromFile(CustomAliasFileName);
    if(err != QERR_SUCCESS) {
      std::stringstream str;
      str<<err;
      Error("%s",str.str().c_str());
    }
  }
  
  std::vector<QEventLabel> evDataLabels = evi.GetWriteLabels();
  std::vector<QEventLabel>::const_iterator evLabelIter = evDataLabels.begin();
  bool inclusive = GetBool("InclusiveMode",true,false);
  while(evLabelIter != evDataLabels.end()) {
    QEventLabel label = *evLabelIter;
    QObject** peventData = evi.Get(label.owner.c_str(),label.name.c_str());
    QObject* eventData = (*peventData);
    
    bool add = !fAppendToInput || eventData->IsJustAdded();
    if(GetBool(label.GetStringLabel(),inclusive,false) && add) {
      std::string newOwner = GetString(label.GetStringLabel() + ":NewOwner","",
                                       false);
      if(newOwner != "") {
        QEventLabel newLabel = label;
        newLabel.owner = newOwner;
        Warn("Changed owner of %s into %s, this is a dangerous action",
             label.GetStringLabel().c_str(),newLabel.GetStringLabel().c_str());
        label = newLabel;
      }
      std::string labelstr = label.GetStringLabel();
      std::string brnamestr = labelstr + ".";
      Debug("Adding branch %s of type %s",
            brnamestr.c_str(),eventData->GetName());
      fTree->Branch(brnamestr.c_str(),eventData->GetName(),peventData);
      
      std::map<std::string, std::string> aliasMap = aliases.Find(label);
      std::map<std::string, std::string>::const_iterator iter =
        aliasMap.begin();
      
      
      while(iter != aliasMap.end()) {
        std::string path;
        if(iter->second =="") {
          path = labelstr;
        } else {
          path = labelstr+std::string(".")+iter->second;
          
        }
        const char* al = iter->first.c_str();
        const char* ap = path.c_str();
        Debug("Adding alias %s for %s",al,ap);
        fTree->SetAlias(al,ap);
        iter++;
      }
    }
    
    evLabelIter++;
  }
  QError err;
  fGlobalWriter =
    QGlobalWriterDispatcher::GetInstance().GetWriter(Q_CURRENT_WRITER,err);
}
//_____________________________________________________________________________
void QRootFileWriterBase::Dump(const QEventInspector& evi) {
  if(fSkipGenerated  && evi.GetEvent().Get<QBool>("Reader","GeneratedEvent"))
    return;
  fTree->Fill();
  if(fGlobalWriter)  
    fGlobalWriter->Set(Form("%s@FlushCache",GetName().c_str()),
                       fRootOutputAddress);
}
//_____________________________________________________________________________
void QRootFileWriterBase::Done() {
  if(fGlobalWriter)
    fGlobalWriter->Set(Form("%s@FlushCache",GetName().c_str()),
                       fRootOutputAddress);
  if(!fAppendToInput) {
    QFileWriter::Done();
    return;
  } 
  
  Close();
 
  AppendToInput();
 
  // remove files
  std::map<std::string,std::string>::iterator iter =
    fReaderWriterFiles.begin();
  while(iter != fReaderWriterFiles.end()) {
    remove(iter->second.c_str());
    iter++;
  }
  RemoveFilesList();
}
void QRootFileWriterBase::AppendToInput() {
  
  // check that number of entries matches
  std::map<std::string,std::string>::const_iterator iter =
    fReaderWriterFiles.begin();
  bool match = true;
  std::string input = "", output = "";
  Long64_t inputEntries = 0, outputEntries = 0;
  while(iter != fReaderWriterFiles.end() && match) {
    input = iter->first;
    output = iter->second;
    Debug("Reader file \"%s\"; writer file \"%s\"",input.c_str(),
          output.c_str());
    
    TFile* ifile = new TFile(input.c_str());
    if(!ifile || ifile->IsZombie() || access(input.c_str(),W_OK)) {
      Error("Input file %s is not writable! Merge with input failed. Input "
            "files not touched. Output saved on specified files.",
            input.c_str());
      SetRunAgain(0);
      WriteFilesList();
      return;
    }
    QTree* itree = (QTree*)ifile->Get("qtree");
    inputEntries = itree->GetEntries();
    TFile* ofile = new TFile(output.c_str());
    QTree* otree = (QTree*)ofile->Get("qtree");
    outputEntries = otree->GetEntries();
    if(inputEntries > 0 && inputEntries == outputEntries) match = true;
    else match = false;
    iter++;
  }
  iter = fReaderWriterFiles.begin();
  if(!match) {
    Error("Input file %s has %llu entries while output file %s has %llu "
          "entries: merge with input failed. Input files not touched. Output"
          " saved on specified files.",input.c_str(),inputEntries,
          output.c_str(),outputEntries);
    SetRunAgain(0);
    WriteFilesList();
    return;
  }
  else {
    while(iter != fReaderWriterFiles.end()){
      // merge on each file
      input = iter->first;
      output = iter->second;
      Info("Merging %s into %s",output.c_str(),input.c_str());
      AppendFileToFile(output.c_str(),input.c_str());
      iter++;
    }
  }
}
//_____________________________________________________________________________
void QRootFileWriterBase::Close() {
  if(fRootOutput) { 
    //File
    std::string fname = fRootOutput->GetName();
    fRootOutput->Write("",TObject::kOverwrite);
    fRootOutput->Close();
    delete fRootOutput;
    fTree = NULL;
    fRootOutput = NULL;
  }
}
//_____________________________________________________________________________
void QRootFileWriterBase::AppendFileToFile(const char *fromFile,
                                           const char *toFile) const {
  // qtree
  TFile ofile(fromFile);
  QTree* otree = (QTree*)ofile.Get("qtree");
  TFile* ifile = new TFile (toFile,"UPDATE");
  QTree* itree = (QTree*)ifile->Get("qtree");
  
  // rename output qtree into qtree_sequenceName and make it
  // friend of qtree.
  TString friendname = "qtree";
  friendname += "_";
  friendname += GetSequence().GetName();
 
  // remove tree in input file
  if(ifile->FindObject(friendname.Data()) && 
          !QOptions::GetInstance().GetBool("General.OverWrite") &&
          !GetBool("OverWriteEvent",false,false)
    ) {
    Panic("AppendToInput: Input file already contains a QTree with name %s. You may want to set the OverWriteEvent=true option to this module or run diana with the option -O or remove the QTree manually with diana-rootfilehandler. Best regards.",friendname.Data());
  }
  ifile->Delete(friendname+";*");
  // remove tree from friend list
  TObject* fele = 0;
  TList* flist = itree->GetListOfFriends();
  if(flist) fele = flist->FindObject(friendname.Data());
  if(fele) itree->GetListOfFriends()->Remove(fele);
  ifile->cd();
  
  // clone new tree and add it to input file
  QTree* otree2 = dynamic_cast<QTree *>(otree->CloneTree(-1,"fast"));
  otree2->SetName(friendname);
  itree->AddFriend(otree2);
  ifile->Write("",TObject::kOverwrite);
 
  // global quantities
  if(TDirectory* dir = ofile.GetDirectory(ROOT_GLOBAL_DIR)) {
    TList* list = dir->GetListOfKeys();
    if(list && list->GetSize() > 0) {
      TDirectory* rootGlobalDir = ifile->GetDirectory(ROOT_GLOBAL_DIR);
      if(rootGlobalDir == 0)
        rootGlobalDir =
          ifile->mkdir(ROOT_GLOBAL_DIR,"Diana global objects");
      rootGlobalDir->cd();
      
      TIter nextobj(list);
      while (TKey* key = (TKey *) nextobj()) {
        TString fullPath = ROOT_GLOBAL_DIR;
        fullPath += "/";
        fullPath += key->GetName();
        TString thisKey = key->GetName();
        rootGlobalDir->Delete(thisKey+";*");
        ofile.Get(fullPath)->Write(key->GetName());
      }
    }
  }
  
  ifile->Write("",TObject::kOverwrite);
  delete ifile;
  
  ofile.Close();
}
//_____________________________________________________________________________