PWGLF/FORWARD/trains/ChainBuilder.C

   1 /**
   2  * @file   ChainBuilder.C
   3  * @author Christian Holm Christensen <cholm@master.hehi.nbi.dk>
   4  * @date   Tue Oct 16 17:54:26 2012
   5  *
   6  * @brief  Build a chain
   7  *
   8  * @ingroup pwglf_forward_trains_util
   9  */
  10
  11 #ifndef CHAINBUILDER_C
  12 #define CHAINBUILDER_C
  13 #ifndef __CINT__
  14 # include <TString.h>
  15 # include <TChain.h>
  16 # include <TSystemDirectory.h>
  17 # include <TSystem.h>
  18 # include <TFile.h>
  19 # include <TList.h>
  20 # include <TError.h>
  21 # include <TROOT.h>
  22 # include <TGridCollection.h>
  23 # include <TFileCollection.h>
  24 # include <THashList.h>
  25 # include <TKey.h>
  26 # include <fstream>
  27 #else
  28 class TString;
  29 class TChain;
  30 class TSystemDirectory;
  31 #endif
  32
  33 // ===================================================================
  34 /**
  35  * Build a chain
  36  *
  37  * @ingroup pwglf_forward_trains_util
  38  */
  39 struct ChainBuilder
  40 {
  41   enum {
  42     kInvalid,
  43     kDirectory,
  44     kXML,
  45     kAscii,
  46     kROOT
  47   };
  48   //------------------------------------------------------------------
  49   static UShort_t CheckSource(TString& src)
  50   {
  51     // Local copy
  52     TString tmp(src);
  53
  54     // --- Normalize the path ----------------------------------------
  55     if (tmp == ".") tmp = "";
  56     if (!tmp.BeginsWith("/")) tmp.Prepend("../");
  57     if (gSystem->ExpandPathName(tmp)) {
  58       Error("ChainBuilder::CheckSource",
  59             "Failed to expand source %s", src.Data());
  60       return kInvalid;
  61     }
  62
  63     // --- Stat the file ---------------------------------------------
  64     FileStat_t stat;
  65     if (gSystem->GetPathInfo(tmp, stat)) return kInvalid;
  66     src = tmp;
  67
  68     // --- Check if directory or file --------------------------------
  69     if (R_ISDIR(stat.fMode)) return kDirectory;
  70
  71     // --- check file type -------------------------------------------
  72     TString type(gSystem->GetFromPipe(Form("file -b %s", src.Data())));
  73     if      (type.Contains("ROOT"))  return kROOT;
  74     else if (type.Contains("XML"))   return kXML;
  75     else if (type.Contains("ASCII")) return kAscii;
  76
  77     Error("ChainBuilder::CheckSource",
  78           "Do not now how to process %s of type %s",
  79           src.Data(), type.Data());
  80     return kInvalid;
  81   }
  82   //------------------------------------------------------------------
  83   /**
  84    * Create the chain.  User is owner.
  85    *
  86    * @return Null in case of problems, chain otherwise
  87    */
  88   static TChain* Create(const TString& src,
  89                         const TString& treeName,
  90                         const TString& pattern,
  91                         Bool_t         mc,
  92                         Bool_t         recursive)
  93   {
  94     TString tmp(src);
  95     UShort_t type = CheckSource(tmp);
  96
  97     return Create(type, tmp, treeName, pattern, mc, recursive);
  98   }
  99   //------------------------------------------------------------------
 100   /**
 101    * Create the chain.  User is owner.
 102    *
 103    * @return Null in case of problems, chain otherwise
 104    */
 105   static TChain* Create(UShort_t       type,
 106                         const TString& src,
 107                         const TString& treeName,
 108                         const TString& pattern,
 109                         Bool_t         mc,
 110                         Bool_t         recursive)
 111   {
 112     // --- check input -----------------------------------------------
 113     if (type == kInvalid) {
 114       Error("ChainBuilder::Create", "Source %s isn't a file or directory",
 115             src.Data());
 116       return 0;
 117     }
 118     TString tN(treeName);
 119     if (tN.IsNull())
 120       Warning("ChainBuilder::Create", "No tree name specified, assuming T");
 121
 122     TString pat(pattern);
 123     if (pat.IsNull()) {
 124       if      (tN.EqualTo("esdTree")) pat = "AliESD";
 125       else if (tN.EqualTo("aodTree")) pat = "AliAOD";
 126     }
 127
 128     // --- Create output ---------------------------------------------
 129     TChain* chain = new TChain(tN);
 130
 131     // --- execute based on type
 132     Bool_t ret = true;
 133     switch (type) {
 134     case kROOT:      ret = CreateFromFile(chain, src); break;
 135     case kXML:       ret = CreateFromXML(chain,  src); break;
 136     case kAscii:     ret = CreateFromList(chain, src); break;
 137     case kDirectory: ret = CreateFromDirectory(chain, src,
 138                                                pat, mc,
 139                                                recursive); break;
 140     default:         ret = false;
 141     }
 142
 143     // --- Clean-up --------------------------------------------------
 144     if (chain->GetListOfFiles()->GetEntries() <= 0) ret = false;
 145     if (!ret) {
 146       delete chain;
 147       chain = 0;
 148     }
 149     return chain;
 150   }
 151   //------------------------------------------------------------------
 152   /**
 153    * Create a chain consiting of a single file
 154    *
 155    * @param fn File name.
 156    *
 157    * @return Chain or null
 158    */
 159   static Bool_t CreateFromFile(TChain* chain, const TString& src)
 160   {
 161     // Info("CreateFromFile", "Making from single file %s", src.Data());
 162     if (!CheckFile(src, chain)) return false;
 163     return true;
 164   }
 165   //------------------------------------------------------------------
 166   /**
 167    * Create a chain from an XML containing an collection
 168    *
 169    * @return Newly allocated chain or null
 170    */
 171   static Bool_t CreateFromXML(TChain* chain, const TString& src)
 172   {
 173     Long_t ret = gROOT->ProcessLine(Form("TAlienCollection(\"%s\")",
 174                                          src.Data()));
 175     if (!ret) {
 176       Error("ChainBuilder::CreateFromXML",
 177             "Cannot create AliEn collection from XML file %s", src.Data());
 178       return false;
 179     }
 180
 181     TGridCollection* collection = reinterpret_cast<TGridCollection*>(ret);
 182     if (!collection) {
 183       Error("ChainBuilder::CreateFromXML",
 184             "Cannot create AliEn collection from XML file %s", src.Data());
 185       return false;
 186     }
 187
 188     collection->Reset();
 189     while (collection->Next()) chain->Add(collection->GetTURL(""));
 190
 191     return true;
 192   }
 193   //------------------------------------------------------------------
 194   /**
 195    * Create a chain from a file containing a list of files
 196    *
 197    * @return Newly allocated chain or null
 198    */
 199   static Bool_t CreateFromList(TChain* chain, const TString& src)
 200   {
 201     std::ifstream in(src.Data());
 202     if (!in) {
 203       Error("ChainBuilder::CreateFromList",
 204             "Failed to open list %s", src.Data());
 205       return false;
 206     }
 207
 208     while (in.good()) {
 209       TString line;
 210       line.ReadToDelim(in);
 211       TString l(line.Strip(TString::kBoth));
 212       if (l.IsWhitespace() || l.BeginsWith("#")) continue;
 213
 214       if (!CheckFile(l, chain))
 215         Warning("ChainBuilder::CreateFromList",
 216                 "Failed to add %s to chain", l.Data());
 217     }
 218     return true;
 219   }
 220   //------------------------------------------------------------------
 221   /**
 222    * Make a chain from a base directory, pattern, and treename -
 223    * possibly recursively
 224    *
 225    * @return true on success
 226    */
 227   static Bool_t CreateFromDirectory(TChain* chain,
 228                                     const TString& src,
 229                                     const TString& pattern,
 230                                     Bool_t         mc,
 231                                     Bool_t         recursive)
 232   {
 233     // Info("", "Scanning src=%s, pattern=%s, mc=%d recursive=%d",
 234     //       src.Data(), pattern.Data(), mc, recursive);
 235     // Save current directory
 236     TString savdir(gSystem->WorkingDirectory());
 237     TSystemDirectory d(gSystem->BaseName(src.Data()), src.Data());
 238     // Info("", "Will scan %s", d.GetTitle());
 239     if (!ScanDirectory(chain, &d, pattern, mc, recursive)) return false;
 240     // Go back to the saved directory
 241     gSystem->ChangeDirectory(savdir);
 242
 243     return true;
 244   }
 245   //------------------------------------------------------------------
 246   /**
 247    * Check if we can add a file to the chain
 248    *
 249    * @param path   Full path to file
 250    * @param chain  Chain
 251    *
 252    * @return true on success, false otherwise
 253    */
 254   static Bool_t CheckFile(const TString& path, TChain* chain)
 255   {
 256     // Info("", "Checking %s", path.Data());
 257     gSystem->RedirectOutput("/dev/null", "w");
 258     TFile* test = TFile::Open(path, "READ");
 259     gSystem->RedirectOutput(0);
 260     if (!test) {
 261       Warning("ChainBuilder::CheckFile", "Failed to open %s", path.Data());
 262       return false;
 263     }
 264
 265     TObject* o = test->Get(chain->GetName());
 266     if (!o) {
 267       // Let's try to find a TFileCollection
 268       TList* l = test->GetListOfKeys();
 269       TIter next(l);
 270       TKey* k = 0;
 271       Bool_t ok = false;
 272       while ((k = static_cast<TKey*>(next()))) {
 273         TString cl(k->GetClassName());
 274         if (!cl.EqualTo("TFileCollection")) continue;
 275         TFileCollection* fc = dynamic_cast<TFileCollection*>(k->ReadObj());
 276         Info("", "Adding file collection");
 277         chain->AddFileInfoList(fc->GetList());
 278         ok = true;
 279       }
 280       if (ok) {
 281         test->Close();
 282         return true;
 283       }
 284     }
 285     else if (dynamic_cast<TTree*>(o)) {
 286       test->Close();
 287       chain->Add(path);
 288       return true;
 289     }
 290
 291     Warning("ChainBuilder::CheckFile",
 292             "The file %s does not contain the tree %s or a file collection",
 293             path.Data(), chain->GetName());
 294
 295     return false;
 296   }
 297   //------------------------------------------------------------------
 298   /**
 299    * Scan directory @a dir (possibly recursive) for tree files to add
 300    * to the chain.    This does not follow sym-links
 301    *
 302    * @param dir        Directory to scan
 303    * @param chain      Chain to add to
 304    *
 305    * @return true if any files where added
 306    */
 307   static Bool_t ScanDirectory(TChain*           chain,
 308                               TSystemDirectory* dir,
 309                               const TString&    pattern,
 310                               Bool_t            mc,
 311                               Bool_t            recursive)
 312   {
 313     // Assume failure
 314     Bool_t ret = false;
 315
 316     // Get list of files, and go back to old working directory
 317     TString oldDir(gSystem->WorkingDirectory());
 318     TList*  files = dir->GetListOfFiles();
 319     if (!gSystem->ChangeDirectory(oldDir)) {
 320       Error("ChainBuilder::ScanDirectory", "Failed to go back to %s",
 321             oldDir.Data());
 322       return false;
 323     }
 324     if (!files) {
 325       Warning("ChainBuilder::ScanDirectory", "No files");
 326       return false;
 327     }
 328
 329     TList toAdd;
 330     toAdd.SetOwner();
 331     Bool_t hasGAlice = (!(mc) ? true : false);
 332     Bool_t hasKine   = (!(mc) ? true : false);
 333     Bool_t hasTrRef  = (!(mc) ? true : false);
 334
 335     // Sort list of files and check if we should add it
 336     files->Sort();
 337     TIter next(files);
 338     TSystemFile* file = 0;
 339     while ((file = static_cast<TSystemFile*>(next()))) {
 340       TString name(file->GetName());
 341       TString title(file->GetTitle());
 342       TString full(gSystem->ConcatFileName(file->GetTitle(), name.Data()));
 343       // Info("", "Got file %s", full.Data());
 344       if (file->IsA()->InheritsFrom(TSystemDirectory::Class())) full = title;
 345       // Ignore special links
 346       if (name == "." || name == "..") {
 347         // Info("ChainBuilder::ScanDirectory", "Ignoring %s", name.Data());
 348         continue;
 349       }
 350       // Info("", "Got file %s", full.Data());
 351
 352       FileStat_t fs;
 353       if (gSystem->GetPathInfo(full.Data(), fs)) {
 354         Warning("ChainBuilder::ScanDirectory", "Cannot stat %s (%s)",
 355                 full.Data(), gSystem->WorkingDirectory());
 356         continue;
 357       }
 358       // Check if this is a directory
 359       if (file->IsDirectory(full)) {
 360         // Info("", "Recursive scan of %s", full.Data());
 361         if (recursive) {
 362           // if (title[0] == '/')
 363           TSystemDirectory* d = new TSystemDirectory(file->GetName(),
 364                                                      full.Data());
 365           if (ScanDirectory(chain, d, pattern, mc, recursive))
 366             ret = true;
 367           delete d;
 368         }
 369         continue;
 370       }
 371
 372       // If this is not a root file, ignore
 373       if (!name.EndsWith(".root")) {
 374         // Info("ScanDirectory", "File %s does not end in .root", name.Data());
 375         continue;
 376       }
 377
 378       // If this file does not contain AliESDs, ignore
 379       if (!name.Contains(pattern)) {
 380         // Info("ChainBuilder::ScanDirectory", "%s does not match pattern %s",
 381         //      name.Data(), pattern.Data());
 382         if (mc) {
 383           if (name.CompareTo("galice.root") == 0)     hasGAlice = true;
 384           if (name.CompareTo("Kinematics.root") == 0) hasKine   = true;
 385           if (name.CompareTo("TrackRefs.root")  == 0) hasTrRef = true;
 386         }
 387         continue;
 388       }
 389
 390       // Add
 391       // Info("ChainBuilder::ScanDirectory", "Adding %s", full.Data());
 392       toAdd.Add(new TObjString(full));
 393     }
 394
 395     if (mc && toAdd.GetEntries() > 0 &&
 396         (!hasGAlice || !hasKine || !hasTrRef)) {
 397       Warning("ChainBuilder::ScanDirectory",
 398               "one or more of {galice,Kinematics,TrackRefs}.root missing from "
 399               "%s, not adding anything from this directory",
 400               dir->GetTitle());
 401       toAdd.Delete();
 402     }
 403
 404     TIter nextAdd(&toAdd);
 405     TObjString* s = 0;
 406     Int_t added = 0;
 407     while ((s = static_cast<TObjString*>(nextAdd()))) {
 408       // Info("ChainBuilder::ScanDirectory",
 409       //      "Adding %s", s->GetString().Data());
 410       TString fn = s->GetString();
 411       if (!CheckFile(fn, chain)) continue;
 412
 413       added++;
 414     }
 415     if (added > 0) ret = true;
 416
 417     gSystem->ChangeDirectory(oldDir);
 418     return ret;
 419   }
 420 };
 421 #endif