#include "TFileMerger.h"
#include "TUrl.h"
#include "TFile.h"
#include "TUUID.h"
#include "TSystem.h"
#include "TH1.h"
#include "TChain.h"
#include "TKey.h"
#include "THashList.h"
#include "TObjString.h"
#include "TClass.h"
#include "TMethodCall.h"
#include "Riostream.h"
ClassImp(TFileMerger)
TFileMerger::TFileMerger(Bool_t isLocal, Bool_t histoonego)
: fOutputFile(0), fFastMethod(kTRUE), fNoTrees(kFALSE),
fLocal(isLocal), fHistoOneGo(histoonego)
{
fFileList = new TList;
fFileList->SetOwner(kTRUE);
fMergeList = new TList;
fMergeList->SetOwner(kTRUE);
}
TFileMerger::~TFileMerger()
{
SafeDelete(fFileList);
SafeDelete(fMergeList);
SafeDelete(fOutputFile);
}
void TFileMerger::Reset()
{
fFileList->Clear();
fMergeList->Clear();
}
Bool_t TFileMerger::AddFile(const char *url)
{
TFile *newfile = 0;
TString localcopy;
if (fLocal) {
TUUID uuid;
localcopy.Form("file:%s/ROOTMERGE-%s.root", gSystem->TempDirectory(), uuid.AsString());
if (!TFile::Cp(url, localcopy)) {
Error("AddFile", "cannot get a local copy of file %s", url);
return kFALSE;
}
newfile = TFile::Open(localcopy, "READ");
} else {
newfile = TFile::Open(url, "READ");
}
if (!newfile) {
if (fLocal)
Error("AddFile", "cannot open local copy %s of URL %s",
localcopy.Data(), url);
else
Error("AddFile", "cannot open file %s", url);
return kFALSE;
} else {
fFileList->Add(newfile);
if (!fMergeList)
fMergeList = new TList;
TObjString *urlObj = new TObjString(url);
fMergeList->Add(urlObj);
return kTRUE;
}
}
Bool_t TFileMerger::OutputFile(const char *outputfile)
{
SafeDelete(fOutputFile);
fOutputFilename = outputfile;
if (!(fOutputFile = TFile::Open(outputfile, "RECREATE"))) {
Error("OutputFile", "cannot open the MERGER output file %s", fOutputFilename.Data());
return kFALSE;
}
return kTRUE;
}
void TFileMerger::PrintFiles(Option_t *options)
{
fFileList->Print(options);
}
Bool_t TFileMerger::Merge(Bool_t)
{
if (!fOutputFile) {
TString outf(fOutputFilename);
if (outf.IsNull()) {
outf.Form("file:%s/FileMerger.root", gSystem->TempDirectory());
Info("Merge", "will merge the results to the file %s\n"
"since you didn't specify a merge filename",
TUrl(outf).GetFile());
}
if (!OutputFile(outf.Data())) {
return kFALSE;
}
}
Bool_t result = MergeRecursive(fOutputFile, fFileList);
if (!result) {
Error("Merge", "error during merge of your ROOT files");
} else {
fOutputFile->Close();
}
SafeDelete(fOutputFile);
TIter next(fFileList);
TFile *file;
while ((file = (TFile*) next())) {
file->Close();
if(fLocal) {
TString p(file->GetPath());
p = p(0, p.Index(':',0));
gSystem->Unlink(p);
}
}
return result;
}
Bool_t TFileMerger::MergeRecursive(TDirectory *target, TList *sourcelist)
{
TString path(target->GetPath());
path.Remove(0, path.Last(':') + 2);
Bool_t addDirStat = TH1::AddDirectoryStatus();
TH1::AddDirectory(kFALSE);
TDirectory *first_source = (TDirectory*)sourcelist->First();
Int_t nguess = sourcelist->GetSize()+1000;
THashList allNames(nguess);
((THashList*)target->GetList())->Rehash(nguess);
((THashList*)target->GetListOfKeys())->Rehash(nguess);
while (first_source) {
TDirectory *current_sourcedir = first_source->GetDirectory(path);
if (!current_sourcedir) {
first_source = (TDirectory*)sourcelist->After(first_source);
continue;
}
TChain *globChain = 0;
TIter nextkey( current_sourcedir->GetListOfKeys() );
TKey *key, *oldkey=0;
while ( (key = (TKey*)nextkey())) {
if (current_sourcedir == target) break;
if (oldkey && !strcmp(oldkey->GetName(),key->GetName())) continue;
if (!strcmp(key->GetClassName(),"TProcessID")) {key->ReadObj(); continue;}
if (allNames.FindObject(key->GetName())) continue;
TClass *cl = TClass::GetClass(key->GetClassName());
if (!cl || !cl->InheritsFrom(TObject::Class())) {
Info("MergeRecursive", "cannot merge object type, name: %s title: %s",
key->GetName(), key->GetTitle());
continue;
}
allNames.Add(new TObjString(key->GetName()));
current_sourcedir->cd();
TObject *obj = key->ReadObj();
if (!obj) {
Info("MergeRecursive", "could not read object for key {%s, %s}",
key->GetName(), key->GetTitle());
continue;
}
if (obj->IsA()->InheritsFrom("TH1")) {
TH1 *h1 = (TH1*)obj;
TList listH;
TFile *nextsource = (TFile*)sourcelist->After( first_source );
while ( nextsource ) {
TDirectory *ndir = nextsource->GetDirectory(path);
if (ndir) {
ndir->cd();
TKey *key2 = (TKey*)gDirectory->GetListOfKeys()->FindObject(key->GetName());
if (key2) {
TObject *hobj = key2->ReadObj();
hobj->ResetBit(kMustCleanup);
listH.Add(hobj);
if (!fHistoOneGo) {
h1->Merge(&listH);
listH.Delete();
}
}
}
nextsource = (TFile*)sourcelist->After( nextsource );
}
if (fHistoOneGo) {
h1->Merge(&listH);
listH.Delete();
}
} else if ( obj->IsA()->InheritsFrom( "TTree" ) ) {
if (!fNoTrees) {
TString obj_name;
if (path.Length()) {
obj_name = path + "/" + obj->GetName();
} else {
obj_name = obj->GetName();
}
globChain = new TChain(obj_name);
globChain->Add(first_source->GetName());
TFile *nextsource = (TFile*)sourcelist->After( first_source );
while ( nextsource ) {
TFile *curf = TFile::Open(nextsource->GetName());
if (curf) {
Bool_t mustAdd = kFALSE;
if (curf->FindKey(obj_name)) {
mustAdd = kTRUE;
} else {
TObject *aobj = curf->Get(obj_name);
if (aobj) { mustAdd = kTRUE; delete aobj;}
}
if (mustAdd) {
globChain->Add(nextsource->GetName());
}
}
delete curf;
nextsource = (TFile*)sourcelist->After( nextsource );
}
}
} else if ( obj->IsA()->InheritsFrom( "TDirectory" ) ) {
target->cd();
TDirectory *newdir = target->mkdir( obj->GetName(), obj->GetTitle() );
MergeRecursive( newdir, sourcelist);
} else if (obj->InheritsFrom(TObject::Class()) &&
obj->IsA()->GetMethodWithPrototype("Merge", "TCollection*") ) {
TList listH;
TString listHargs;
listHargs.Form("((TCollection*)0x%lx)",&listH);
TFile *nextsource = (TFile*)sourcelist->After( first_source );
while (nextsource) {
TDirectory *ndir = nextsource->GetDirectory(path);
if (ndir) {
ndir->cd();
TKey *key2 = (TKey*)gDirectory->GetListOfKeys()->FindObject(key->GetName());
if (key2) {
TObject *hobj = key2->ReadObj();
if (hobj->InheritsFrom(TCollection::Class())) {
((TCollection*)hobj)->SetOwner();
}
hobj->ResetBit(kMustCleanup);
listH.Add(hobj);
Int_t error = 0;
obj->Execute("Merge", listHargs.Data(), &error);
if (error) {
Error("MergeRecursive", "calling Merge() on '%s' with the corresponding object in '%s'",
obj->GetName(), nextsource->GetName());
}
listH.Delete();
}
}
nextsource = (TFile*)sourcelist->After( nextsource );
}
} else {
Warning("MergeRecursive", "cannot merge object type (n:'%s', t:'%s') - "
"Merge(TCollection *) not implemented",
obj->GetName(), obj->GetTitle());
TFile *nextsource = (TFile*)sourcelist->After( first_source );
while (nextsource) {
TDirectory *ndir = nextsource->GetDirectory(path);
if (ndir) {
ndir->cd();
TKey *key2 = (TKey*)gDirectory->GetListOfKeys()->FindObject(key->GetName());
if (key2) {
TObject *nobj = key2->ReadObj();
nobj->ResetBit(kMustCleanup);
if (target->WriteTObject(nobj, key2->GetName(), "SingleKey") <= 0) {
Warning("MergeRecursive", "problems copying object (n:'%s', t:'%s') to output file ",
obj->GetName(), obj->GetTitle());
}
delete nobj;
}
}
nextsource = (TFile*)sourcelist->After( nextsource );
}
}
target->cd();
if(obj->IsA()->InheritsFrom( "TDirectory" )) {
} else if(obj->IsA()->InheritsFrom( "TTree" )) {
if (!fNoTrees) {
if (globChain) {
globChain->ls();
if (fFastMethod) globChain->Merge(target->GetFile(),0,"keep fast");
else globChain->Merge(target->GetFile(),0,"keep");
delete globChain;
}
}
} else if (obj->IsA()->InheritsFrom( "TCollection" )) {
obj->Write( key->GetName(), TObject::kSingleKey );
((TCollection*)obj)->SetOwner();
} else {
obj->Write( key->GetName() );
}
if (obj->IsA()->InheritsFrom("TCollection")) ((TCollection*)obj)->Delete();
oldkey = key;
delete obj;
}
first_source = (TDirectory*)sourcelist->After(first_source);
}
target->SaveSelf(kTRUE);
TH1::AddDirectory(addDirStat);
return kTRUE;
}