ToolsDiscovery.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
*/

 
//
// includes
//
 
#include <SampleHandler/ToolsDiscovery.h>
 
#include <RootCoreUtils/Assert.h>
#include <RootCoreUtils/StringUtil.h>
#include <SampleHandler/DiskListLocal.h>
#include <SampleHandler/GridTools.h>
#include <SampleHandler/MessageCheck.h>
#include <SampleHandler/MetaFields.h>
#include <SampleHandler/MetaObject.h>
#include <SampleHandler/SampleGrid.h>
#include <SampleHandler/SampleHandler.h>
#include <SampleHandler/SampleLocal.h>
#include <SampleHandler/ScanDir.h>
#include <TChain.h>
#include <TChainElement.h>
#include <TFile.h>
#include <TTree.h>
#include <fstream>
#include <memory>
#include <stdexcept>
 
//
// method implementations
//
 
namespace SH
{
  void scanDir (SampleHandler& sh, DiskList& list,
        const std::string& pattern,
        const std::string& samplePattern,
        const std::string& samplePostfix)
  {
    ScanDir()
      .sampleDepth (0)
      .minDepth (1)
      .maxDepth (1)
      .filePattern (pattern)
      .samplePattern (samplePattern)
      .samplePostfix (samplePostfix)
      .scan (sh, list);
  }
 
 
 
  void scanDir (SampleHandler& sh, const std::string& dir)
  {
    ScanDir()
      .sampleDepth (0)
      .minDepth (1)
      .maxDepth (1)
      .scan (sh, dir);
  }
 
 
 
  void scanDir (SampleHandler& sh, const std::string& dir,
        const std::string& prefix)
  {
    DiskListLocal list (dir, prefix);
    ScanDir()
      .sampleDepth (0)
      .minDepth (1)
      .maxDepth (1)
      .scan (sh, list);
  }
 
 
 
  void scanFiles (SampleHandler& sh, DiskList& list,
          const std::string& pattern)
  {
    ScanDir()
      .sampleDepth (-1)
      .filePattern (pattern)
      .samplePostfix (".*")
      .scan (sh, list);
  }
 
 
 
  std::unique_ptr<Sample> makeFromTChainUnique (const std::string& name, const TChain& chain)
  {
    auto result = std::make_unique<SampleLocal> (name);
    result->meta()->setString (MetaFields::treeName, chain.GetName());
 
    TIter chainIter (chain.GetListOfFiles());
    TChainElement *chainElement = nullptr;
    while ((chainElement = dynamic_cast<TChainElement*>(chainIter.Next())) != nullptr)
      result->add (chainElement->GetTitle());
    return result;
  }
 
 
 
  Sample *makeFromTChain (const std::string& name, const TChain& chain)
  {
    return makeFromTChainUnique (name, chain).release();
  }
 
 
 
  void scanSingleDir (SampleHandler& sh, const std::string& name,
              DiskList& list, const std::string& pattern)
  {
    ScanDir()
      .sampleDepth (0)
      .filePattern (pattern)
      .sampleRename ("*", name)
      .scan (sh, list);
  }
 
 
 
  void scanDQ2 (SampleHandler& sh, const std::string& pattern)
  {
    if (pattern.find ("*") == std::string::npos)
    {
      addGrid (sh, pattern);
    } else
    {
      std::set<std::string> types = {"DATASET", "DIDType.DATASET"};
      if (!pattern.empty() && pattern.back() == '/')
    types = {"CONTAINER", "DIDType.CONTAINER"};
 
      auto subresult = rucioListDids (pattern);
      for (auto& ds : subresult)
      {
    if (types.find (ds.type) != types.end())
      addGrid (sh, ds.scope + ":" + ds.name);
      }
    }
  }
 
 
 
  void scanRucio (SampleHandler& sh, const std::string& pattern,
          bool alwaysQuery)
  {
    if (pattern.find ("*") == std::string::npos && !alwaysQuery)
    {
      addGrid (sh, pattern);
    } else
    {
      auto subresult = rucioListDids (pattern);
      bool added = false;
      for (std::string type : {"CONTAINER", "DIDType.CONTAINER", "DATASET", "DIDType.DATASET"})
      {
    for (auto& ds : subresult)
    {
      if (ds.type == type)
      {
        addGrid (sh, ds.scope + ":" + ds.name);
        added = true;
      }
    }
    if (added)
      return;
      }
      throw std::runtime_error ("failed to find any datasets matching pattern: " + pattern);
    }
  }
 
 
 
  void addGrid (SampleHandler& sh, const std::string& ds)
  {
    RCU_ASSERT_SOFT (ds.find ("*") == std::string::npos);
 
    std::string name;
    if (!ds.empty() && ds.back() == '/')
      name = ds.substr (0, ds.size()-1);
    else
      name = ds;
 
    auto sample = std::make_unique<SampleGrid> (name);
    sample->meta()->setString (MetaFields::gridName, ds);
    sample->meta()->setString (MetaFields::gridFilter, MetaFields::gridFilter_default);
    sh.add (std::move (sample));
  }
 
 
 
  void addGridCombined (SampleHandler& sh, const std::string& dsName,
                        const std::vector<std::string>& dsList)
  {
    std::string name;
    for (const std::string &ds : dsList)
    {
      RCU_ASSERT_SOFT (ds.find ("*") == std::string::npos);
 
      if (!name.empty())
        name.append(",");
 
      if (ds.at(ds.size() - 1) == '/')
        name.append(ds.substr (0, ds.size() - 1));
      else
        name.append(ds);
    }
 
    auto sample = std::make_unique<SampleGrid> (dsName);
    sample->meta()->setString (MetaFields::gridName, name);
    sample->meta()->setString (MetaFields::gridFilter, MetaFields::gridFilter_default);
    sh.add (std::move (sample));
  }
  void addGridCombinedFromFile (SampleHandler& sh, const std::string& dsName,
                                const std::string& dsFile)
  {
    std::ifstream file (dsFile.c_str());
 
    std::string name;
    std::string ds;
    const std::set<char> whitespaces{'\t',' ','\n','\r'};
    while (std::getline (file, ds))
    {
      while ((!ds.empty()) && whitespaces.count(ds.back())) ds.pop_back();
      if (ds.empty() || ds.at(0) == '#')
        continue;
 
      RCU_ASSERT_SOFT (ds.find ("*") == std::string::npos);
 
      if (!name.empty())
        name.append(",");
 
      if (ds.at(ds.size() - 1) == '/')
        name.append(ds.substr (0, ds.size() - 1));
      else
        name.append(ds);
    }
    if (!file.eof())
      throw std::runtime_error ("failed to read file: " + dsFile);
 
    auto sample = std::make_unique<SampleGrid> (dsName);
    sample->meta()->setString (MetaFields::gridName, name);
    sample->meta()->setString (MetaFields::gridFilter, MetaFields::gridFilter_default);
    sh.add (std::move (sample));
  }
 
 
  void makeGridDirect (SampleHandler& sh, const std::string& disk,
               const std::string& from, const std::string& to,
               bool allow_partial)
  {
    using namespace msgDiscovery;
 
    SampleHandler mysh;
 
    for (auto sample : sh.samples())
    {
      SampleGrid *grid = dynamic_cast<SampleGrid*>(sample.get());
 
      if (grid == nullptr)
      {
        mysh.add (sample);
      } else
      {
        const std::string ds = grid->meta()->castString (MetaFields::gridName);
        if (ds.empty())
          throw std::runtime_error ("no dataset configured for grid sample " + grid->name());
 
        std::regex pattern (RCU::glob_to_regexp (grid->meta()->castString (MetaFields::gridFilter, MetaFields::gridFilter_default)));
 
        std::set<std::string> knownFiles;
        std::map<std::string,std::string> usedFiles;
        for (auto& entry : rucioListFileReplicas (ds))
        {
          if (RCU::match_expr (pattern, entry.name))
          {
            knownFiles.insert (entry.name);
            if (entry.disk == disk)
            {
              std::string url = entry.replica;
              const auto split = url.find (from);
              if (split == std::string::npos)
                throw std::runtime_error ("prefix \"" + from + "\" not part of replica URL: " + url);
              url.replace(split, from.size(), to);
              usedFiles[entry.name] = url;
            }
          }
        }
 
        if (usedFiles.empty())
        {
          if (allow_partial)
            ANA_MSG_WARNING ("dataset " << ds << " not at " << disk << ", skipped");
        } else if (knownFiles.size() != usedFiles.size())
        {
          if (allow_partial)
          {
            ANA_MSG_WARNING ("only incomplete version of dataset " << ds << " at " << disk);
          } else
          {
            usedFiles.clear ();
          }
        }
 
        if (usedFiles.size() == 0)
        {
          mysh.add (sample);
        } else
        {
          auto mysample = std::make_unique<SampleLocal> (grid->name());
          *mysample->meta() = *grid->meta();
 
          for (const auto& file : usedFiles)
          {
            mysample->add (file.second);
          }
          mysh.add (std::move (mysample));
        }
      }
    }
    swap (sh, mysh);
  }
 
 
 
  void scanForTrees (SampleHandler& sh, const std::shared_ptr<Sample>& sample,
             const std::string& pattern)
  {
    auto mysample = sample->makeLocal();
    if (mysample->numFiles() == 0)
    {
      sh.add (sample);
      return;
    }
    std::unique_ptr<TFile> file (TFile::Open (mysample->fileName(0).c_str()));
    if (!file.get())
      throw std::runtime_error ("could not open file: " + mysample->fileName(0));
    TObject *object = nullptr;
    std::regex mypattern (pattern);
    // rationale: GetListOfKeys() contains one key per tree cycle
    //   (e.g. "physics;1", "physics;2"), all with the same name.  we
    //   process each name only once, so that repeated cycles do not
    //   make sh.add throw on a duplicate sample name.
    std::set<std::string> seenTrees;
    for (TIter iter (file->GetListOfKeys()); (object = iter.Next()); )
    {
      if (RCU::match_expr (mypattern, object->GetName()) &&
      seenTrees.insert (object->GetName()).second &&
      dynamic_cast<TTree*>(file->Get(object->GetName())))
      {
    std::string newName = sample->name() + "_" + object->GetName();
        std::unique_ptr<Sample> newSample
      (dynamic_cast<Sample*>(sample->Clone (newName.c_str())));
        newSample->name (newName);
        newSample->meta()->setString (MetaFields::treeName, object->GetName());
    sh.add (std::move (newSample));
      }
    }
  }
 
 
 
  void scanForTrees (SampleHandler& sh, const std::string& pattern)
  {
    SH::SampleHandler sh_new;
 
    for (auto sample : sh.samples())
    {
      scanForTrees (sh_new, sample, pattern);
    }
    swap (sh, sh_new);
  }
 
 
 
  void readFileList (SampleHandler& sh, const std::string& name,
             const std::string& file)
  {
    std::ifstream myfile (file.c_str());
 
    auto sample = std::make_unique<SampleLocal> (name);
    std::string line;
    const std::set<char> whitespaces{'\t',' ','\n','\r'};
    while (std::getline (myfile, line))
    {
      while ((!line.empty()) && whitespaces.count(line.back())) line.pop_back();
      if (!line.empty() && line.at(0) != '#')
      {
    sample->add (line);
      }
    }
    if (!myfile.eof())
      throw std::runtime_error ("failed to read file: " + file);
    sh.add (std::move (sample));
  }
}