[rfile] add GetKeys() methods

Author: silverweed

io/io/inc/ROOT/RFile.hxx

@@ -11,10 +11,12 @@
 #include <ROOT/RError.hxx>
 
 #include <memory>
+#include <iostream>
 #include <string_view>
 #include <typeinfo>
 
 class TFile;
+class TIterator;
 class TKey;
 
 namespace ROOT {
@@ -29,6 +31,95 @@ ROOT::RLogChannel &RFileLog();
 
 } // namespace Internal
 
+/**
+\class ROOT::Experimental::RKeyInfo
+\ingroup RFile
+\brief Information about an RFile object's Key.
+
+Every object inside a ROOT file has an associated "Key" which contains metadata on the object, such as its name, type
+etc.
+Querying this information can be done via RFile::ListKeys(). Reading an object's Key
+doesn't deserialize the full object, so it's a relatively lightweight operation.
+*/
+struct RKeyInfo {
+   enum class ECategory : std::uint16_t {
+      kInvalid,
+      kObject,
+      kDirectory
+   };
+
+   std::string fName;
+   std::string fTitle;
+   std::string fClassName;
+   std::uint16_t fCycle = 0;
+   ECategory fCategory = ECategory::kInvalid;
+};
+
+/// The iterable returned by RFile::ListKeys()
+class RFileKeyIterable final {
+   using Pattern_t = std::string;
+
+   TFile *fFile = nullptr;
+   Pattern_t fPattern;
+   std::uint32_t fFlags = 0;
+
+public:
+   class RIterator {
+      friend class RFileKeyIterable;
+
+      struct RIterStackElem {
+         // This is ugly, but TList returns an (owning) pointer to a polymorphic TIterator...and we need this class
+         // to be copy-constructible.
+         std::shared_ptr<TIterator> fIter;
+         std::string fDirPath;
+
+         // Outlined to avoid including TIterator.h
+         RIterStackElem(TIterator *it, const std::string &path = "");
+         // Outlined to avoid including TIterator.h
+         ~RIterStackElem();
+
+         // fDirPath doesn't need to be compared because it's implied by fIter.
+         bool operator==(const RIterStackElem &other) const { return fIter == other.fIter; }
+      };
+
+      std::vector<RIterStackElem> fIterStack;
+      Pattern_t fPattern;
+      const TKey *fCurKey = nullptr;
+      std::uint16_t fRootDirNesting = 0;
+      std::uint32_t fFlags = 0;
+
+      void Advance();
+
+      // NOTE: `iter` here is an owning pointer (or null)
+      RIterator(TIterator *iter, Pattern_t pattern, std::uint32_t flags);
+
+   public:
+      using iterator = RIterator;
+      using iterator_category = std::forward_iterator_tag;
+      using difference_type = std::ptrdiff_t;
+      using value_type = RKeyInfo;
+      using pointer = const value_type *;
+      using reference = const value_type &;
+
+      iterator &operator++()
+      {
+         Advance();
+         return *this;
+      }
+      value_type operator*();
+      bool operator!=(const iterator &rh) const { return !(*this == rh); }
+      bool operator==(const iterator &rh) const { return fIterStack == rh.fIterStack; }
+   };
+
+   RFileKeyIterable(TFile *file, std::string_view rootDir, std::uint32_t flags)
+      : fFile(file), fPattern(std::string(rootDir)), fFlags(flags)
+   {
+   }
+
+   RIterator begin() const;
+   RIterator end() const;
+};
+
 /**
 \class ROOT::Experimental::RFile
 \ingroup RFile
@@ -126,6 +217,12 @@ class RFile final {
    TKey *GetTKey(std::string_view path) const;
 
 public:
+   enum EListKeyFlags {
+      kListObjects = 1 << 0,
+      kListDirs = 1 << 1,
+      kListRecursive = 1 << 2,
+   };
+
    // This is arbitrary, but it's useful to avoid pathological cases
    static constexpr int kMaxPathNesting = 1000;
 
@@ -196,6 +293,23 @@ public:
 
    /// Flushes the RFile if needed and closes it, disallowing any further reading or writing.
    void Close();
+
+   /// Returns an iterable over all paths of objects written into this RFile starting at path "rootPath".
+   /// The returned paths are always "absolute" paths: they are not relative to `rootPath`.
+   /// Keys relative to directories are not returned: only those relative to leaf objects are.
+   /// If `rootPath` is the path of a leaf object, only `rootPath` itself will be returned.
+   /// `flags` is a bitmask specifying the listing mode.
+   /// If `(flags & kListObject) != 0`, the listing will include keys of non-directory objects (default);
+   /// If `(flags & kListDirs) != 0`, the listing will include keys of directory objects;
+   /// If `(flags & kListRecursive) != 0`, the listing will recurse on all subdirectories of `rootPath` (default),
+   /// otherwise it will only list immediate children of `rootPath`.
+   RFileKeyIterable ListKeys(std::string_view rootPath = "", std::uint32_t flags = kListObjects | kListRecursive) const
+   {
+      return RFileKeyIterable(fFile.get(), rootPath, flags);
+   }
+
+   /// Prints the internal structure of this RFile to the given stream.
+   void Print(std::ostream &out = std::cout) const;
 };
 
 } // namespace Experimental

io/io/src/RFile.cxx

@@ -13,7 +13,9 @@
 #include <Byteswap.h>
 #include <TError.h>
 #include <TFile.h>
+#include <TIterator.h>
 #include <TKey.h>
+#include <TList.h>
 #include <TROOT.h>
 
 #include <algorithm>
@@ -169,18 +171,14 @@ static std::string ValidateAndNormalizePath(std::string &path)
 
 /////////////////////////////////////////////////////////////////////////////////////////////////
 
-RFile::RFile(std::unique_ptr<TFile> file) : fFile(std::move(file)) {}
-
-RFile::~RFile() = default;
-
 std::unique_ptr<RFile> RFile::Open(std::string_view path)
 {
    TDirectory::TContext ctx(nullptr); // XXX: probably not thread safe?
    auto tfile = std::unique_ptr<TFile>(TFile::Open(std::string(path).c_str(), "READ_WITHOUT_GLOBALREGISTRATION"));
    if (!tfile || tfile->IsZombie())
       throw ROOT::RException(R__FAIL("failed to open file " + std::string(path) + " for reading"));
 
-   if (tfile->IsRaw())
+   if (tfile->IsRaw() || !tfile->IsBinary())
       throw ROOT::RException(R__FAIL("Opened file " + std::string(path) + " is not a ROOT file"));
 
    auto rfile = std::unique_ptr<RFile>(new RFile(std::move(tfile)));
@@ -194,7 +192,7 @@ std::unique_ptr<RFile> RFile::Update(std::string_view path)
    if (!tfile || tfile->IsZombie())
       throw ROOT::RException(R__FAIL("failed to open file " + std::string(path) + " for updating"));
 
-   if (tfile->IsRaw())
+   if (tfile->IsRaw() || !tfile->IsBinary())
       throw ROOT::RException(R__FAIL("Opened file " + std::string(path) + " is not a ROOT file"));
 
    auto rfile = std::unique_ptr<RFile>(new RFile(std::move(tfile)));
@@ -208,13 +206,17 @@ std::unique_ptr<RFile> RFile::Recreate(std::string_view path)
    if (!tfile || tfile->IsZombie())
       throw ROOT::RException(R__FAIL("failed to open file " + std::string(path) + " for writing"));
 
-   if (tfile->IsRaw())
+   if (tfile->IsRaw() || !tfile->IsBinary())
       throw ROOT::RException(R__FAIL("Opened file " + std::string(path) + " is not a ROOT file"));
 
    auto rfile = std::unique_ptr<RFile>(new RFile(std::move(tfile)));
    return rfile;
 }
 
+RFile::RFile(std::unique_ptr<TFile> file) : fFile(std::move(file)) {}
+
+RFile::~RFile() = default;
+
 TKey *RFile::GetTKey(std::string_view path) const
 {
    // In RFile, differently from TFile, when dealing with a path like "a/b/c", we always consider it to mean
@@ -361,6 +363,138 @@ void RFile::PutUntyped(std::string_view pathSV, const std::type_info &type, cons
    }
 }
 
+ROOT::Experimental::RFileKeyIterable::RIterator::RIterStackElem::RIterStackElem(TIterator *it, const std::string &path)
+   : fIter(it), fDirPath(path)
+{
+}
+
+ROOT::Experimental::RFileKeyIterable::RIterator::RIterStackElem::~RIterStackElem() = default;
+
+ROOT::Experimental::RFileKeyIterable::RIterator::RIterator(TIterator *iter, Pattern_t pattern, std::uint32_t flags)
+   : fPattern(pattern), fFlags(flags)
+{
+   if (iter) {
+      fIterStack.emplace_back(iter);
+
+      if (!pattern.empty()) {
+         fRootDirNesting = std::count(pattern.begin(), pattern.end(), '/');
+         // `pattern` may or may not end with '/', but we consider it a directory regardless.
+         // In other words, like in virtually all filesystem operations, "dir" and "dir/" are equivalent.
+         fRootDirNesting += pattern.back() != '/';
+      }
+
+      // Advance the iterator to skip the first key, which is always the TFile key.
+      // This will also skip keys until we reach the first correct key we want to return.
+      Advance();
+   }
+}
+
+ROOT::Experimental::RFileKeyIterable::RIterator ROOT::Experimental::RFileKeyIterable::begin() const
+{
+   return {fFile->GetListOfKeys()->MakeIterator(), fPattern, fFlags};
+}
+
+ROOT::Experimental::RFileKeyIterable::RIterator ROOT::Experimental::RFileKeyIterable::end() const
+{
+   return {nullptr, fPattern, fFlags};
+}
+
+void ROOT::Experimental::RFileKeyIterable::RIterator::Advance()
+{
+   fCurKey = nullptr;
+
+   const bool recursive = fFlags & RFile::kListRecursive;
+   const bool includeObj = fFlags & RFile::kListObjects;
+   const bool includeDirs = fFlags & RFile::kListDirs;
+
+   // We only want to return keys that refer to user objects, not internal ones, therefore we skip
+   // all keys that have internal class names.
+   while (!fIterStack.empty()) {
+      auto &[iter, dirPath] = fIterStack.back();
+      assert(iter);
+      TObject *keyObj = iter->Next();
+      if (!keyObj) {
+         // reached end of the iteration
+         fIterStack.pop_back();
+         continue;
+      }
+
+      assert(keyObj->IsA() == TClass::GetClass<TKey>());
+      auto key = static_cast<TKey *>(keyObj);
+
+      const auto dirSep = (dirPath.empty() ? "" : "/");
+      const auto isDir =
+         strcmp(key->GetClassName(), "TDirectory") == 0 || strcmp(key->GetClassName(), "TDirectoryFile") == 0;
+
+      if (isDir) {
+         TDirectory *dir = key->ReadObject<TDirectory>();
+         TIterator *innerIter = dir->GetListOfKeys()->MakeIterator();
+         assert(innerIter);
+         fIterStack.emplace_back(innerIter, dirPath + dirSep + dir->GetName());
+         if (!includeDirs)
+            continue;
+      } else if (!includeObj) {
+         continue;
+      }
+
+      // Reconstruct the full path of the key
+      const auto &fullPath = dirPath + dirSep + key->GetName();
+      const auto nesting = fIterStack.size() - 1;
+
+      // Skip key if it's not a child of root dir
+      if (!ROOT::StartsWith(fullPath, fPattern))
+         continue;
+
+      // Check that we are in the same directory as "rootDir".
+      // Note that for directories we list both the root dir and the immediate children (in non-recursive mode).
+      if (!recursive && nesting != fRootDirNesting && (!isDir || nesting != fRootDirNesting + 1))
+         continue;
+
+      // All checks passed: return this key.
+      assert(!fullPath.empty());
+      fCurKey = key;
+      break;
+   }
+}
+
+ROOT::Experimental::RKeyInfo ROOT::Experimental::RFileKeyIterable::RIterator::operator*()
+{
+   if (fIterStack.empty())
+      throw ROOT::RException(R__FAIL("tried to dereference an invalid iterator"));
+
+   const TKey *key = fCurKey;
+   if (!key)
+      throw ROOT::RException(R__FAIL("tried to dereference an invalid iterator"));
+
+   const bool isDir =
+      strcmp(key->GetClassName(), "TDirectory") == 0 || strcmp(key->GetClassName(), "TDirectoryFile") == 0;
+   const auto &dirPath = fIterStack.back().fDirPath;
+
+   RKeyInfo keyInfo;
+   keyInfo.fCategory = isDir ? RKeyInfo::ECategory::kDirectory : RKeyInfo::ECategory::kObject;
+   keyInfo.fName = dirPath;
+   if (!isDir)
+      keyInfo.fName += std::string(dirPath.empty() ? "" : "/") + key->GetName();
+   keyInfo.fClassName = key->GetClassName();
+   keyInfo.fCycle = key->GetCycle();
+   keyInfo.fTitle = key->GetTitle();
+   return keyInfo;
+}
+
+void RFile::Print(std::ostream &out) const
+{
+   std::vector<RKeyInfo> keys;
+   auto keysIter = ListKeys();
+   for (const auto &key : keysIter) {
+      keys.emplace_back(key);
+   }
+
+   std::sort(keys.begin(), keys.end(), [](const auto &a, const auto &b) { return a.fName < b.fName; });
+   for (const auto &key : keys) {
+      out << key.fClassName << " " << key.fName << ";" << key.fCycle << ": \"" << key.fTitle << "\"\n";
+   }
+}
+
 size_t RFile::Flush()
 {
    return fFile->Write();