[rfile] Add pythonizations

bindings/pyroot/pythonizations/CMakeLists.txt

@@ -86,6 +86,7 @@ set(py_sources
   ROOT/_pythonization/_generic.py
   ROOT/_pythonization/_memory_utils.py
   ROOT/_pythonization/_pyz_utils.py
+  ROOT/_pythonization/_rfile.py
   ROOT/_pythonization/_rntuple.py
   ROOT/_pythonization/_runtime_error.py
   ROOT/_pythonization/_rvec.py

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py

@@ -0,0 +1,148 @@
+# Author: Giacomo Parolini CERN 04/2025
+
+r"""
+\pythondoc RFile
+
+TODO: document RFile
+
+\code{.py}
+# TODO code example
+\endcode
+
+\endpythondoc
+"""
+
+from . import pythonization
+
+
+class _RFile_Get:
+    """
+    Allow access to objects through the method Get().
+    This is pythonized to allow Get() to be called both with and without a template argument.
+    """
+
+    def __init__(self, rfile):
+        self._rfile = rfile
+
+    def __call__(self, namecycle):
+        """
+        Non-templated Get()
+        """
+        import ROOT
+        import cppyy
+
+        key = self._rfile.GetKeyInfo(namecycle)
+        if key:
+            obj = ROOT.Experimental.Internal.RFile_GetObjectFromKey(self._rfile, key)
+            return cppyy.bind_object(obj, key.GetClassName())
+        # No key
+        return None
+
+    def __getitem__(self, template_arg):
+        """
+        Templated Get()
+        """
+
+        def getitem_wrapper(namecycle):
+            obj = self._rfile._OriginalGet[template_arg](namecycle)
+            return obj if obj else None
+
+        return getitem_wrapper
+
+
+class _RFile_Put:
+    """
+    Allow writing objects through the method Put().
+    This is pythonized to allow Put() to be called both with and without a template argument.
+    """
+
+    def __init__(self, rfile):
+        self._rfile = rfile
+
+    def __call__(self, name, obj):
+        """
+        Non-templated Put()
+        """
+        objType = type(obj)
+        if objType == str:
+            # special case: automatically convert python str to std::string
+            className = "std::string"
+        elif not hasattr(objType, '__cpp_name__'):
+            raise TypeError(f"type {objType} is not supported by ROOT I/O") 
+        else:
+            className = objType.__cpp_name__
+        self._rfile.Put[className](name, obj)
+
+    def __getitem__(self, template_arg):
+        """
+        Templated Put()
+        """
+        return self._rfile._OriginalPut[template_arg]
+
+
+def _RFileExit(obj, exc_type, exc_val, exc_tb):
+    """
+    Close the RFile object.
+    Signature and return value are imposed by Python, see
+    https://docs.python.org/3/library/stdtypes.html#typecontextmanager.
+    """
+    obj.Close()
+    return False
+
+
+def _RFileOpen(original):
+    """
+    Pythonization for the factory methods (Recreate, Open, Update)
+    """
+
+    def rfile_open_wrapper(klass, *args):
+        rfile = original(*args)
+        rfile._OriginalGet = rfile.Get
+        rfile.Get = _RFile_Get(rfile)
+        rfile._OriginalPut = rfile.Put
+        rfile.Put = _RFile_Put(rfile)
+        return rfile
+
+    return rfile_open_wrapper
+
+
+def _RFileInit(rfile):
+    """
+    Prevent the creation of RFile through constructor (must use a factory method)
+    """
+    raise NotImplementedError("RFile can only be created via Recreate, Open or Update")
+
+
+def _GetKeyInfo(rfile, path):
+    key = rfile._OriginalGetKeyInfo(path)
+    if key.has_value():
+        return key.value()
+    return None
+
+
+def _ListKeys(rfile, basePath="", listObjects = True, listDirs = False, listRecursive = True):
+    from ROOT.Experimental import RFile
+
+    flags = (listObjects * RFile.kListObjects) | (listDirs * RFile.kListDirs) | (listRecursive * RFile.kListRecursive)
+    iter = rfile._OriginalListKeys(basePath, flags)
+    return iter
+
+
+@pythonization("RFile", ns="ROOT::Experimental")
+def pythonize_rfile(klass):
+    # Explicitly prevent to create a RFile via ctor
+    klass.__init__ = _RFileInit
+
+    # Pythonize factory methods
+    klass.Open = classmethod(_RFileOpen(klass.Open))
+    klass.Update = classmethod(_RFileOpen(klass.Update))
+    klass.Recreate = classmethod(_RFileOpen(klass.Recreate))
+
+    # Pythonization for __enter__ and __exit__ methods
+    # These make RFile usable in a `with` statement as a context manager
+    klass.__enter__ = lambda rfile: rfile
+    klass.__exit__ = _RFileExit
+    klass._OriginalGetKeyInfo = klass.GetKeyInfo
+    klass.GetKeyInfo = _GetKeyInfo
+    klass._OriginalListKeys = klass.ListKeys
+    klass.ListKeys = _ListKeys

io/io/inc/ROOT/RFile.hxx

@@ -11,10 +11,12 @@
 #include <ROOT/RError.hxx>
 
 #include <deque>
-#include <memory>
+#include <functional>
 #include <iostream>
+#include <memory>
 #include <string_view>
 #include <typeinfo>
+#include <variant>
 
 class TFile;
 class TIterator;
@@ -23,13 +25,17 @@ class TKey;
 namespace ROOT {
 namespace Experimental {
 
+class RKeyInfo;
 class RFile;
-struct RFileKeyInfo;
 
 namespace Internal {
 
 ROOT::RLogChannel &RFileLog();
 
+/// Returns an **owning** pointer to the object referenced by `key`. The caller must delete this pointer.
+/// This method is meant to only be used by the pythonization.
+[[nodiscard]] void *RFile_GetObjectFromKey(RFile &file, const RKeyInfo &key);
+
 } // namespace Internal
 
 namespace Detail {
@@ -59,6 +65,7 @@ Querying this information can be done via RFile::ListKeys(). Reading an object's
 doesn't deserialize the full object, so it's a relatively lightweight operation.
 */
 class RKeyInfo final {
+   friend class ROOT::Experimental::RFile;
    friend class ROOT::Experimental::RFileKeyIterable;
 
 public:
@@ -216,6 +223,8 @@ auto myObj = file->Get<TH1D>("h");
 ~~~
 */
 class RFile final {
+   friend void *Internal::RFile_GetObjectFromKey(RFile &file, const RKeyInfo &key);
+
    /// Flags used in PutInternal()
    enum PutFlags {
       /// When encountering an object at the specified path, overwrite it with the new one instead of erroring out.
@@ -231,7 +240,8 @@ class RFile final {
 
    /// Gets object `path` from the file and returns an **owning** pointer to it.
    /// The caller should immediately wrap it into a unique_ptr of the type described by `type`.
-   [[nodiscard]] void *GetUntyped(std::string_view path, const std::type_info &type) const;
+   [[nodiscard]] void *GetUntyped(std::string_view path,
+                                  std::variant<const char *, std::reference_wrapper<const std::type_info>> type) const;
 
    /// Writes `obj` to file, without taking its ownership.
    void PutUntyped(std::string_view path, const std::type_info &type, const void *obj, std::uint32_t flags);
@@ -357,6 +367,9 @@ public:
       return RFileKeyIterable(fFile.get(), basePath, flags);
    }
 
+   /// Retrieves information about the key of object at `path`, if one exists.
+   std::optional<RKeyInfo> GetKeyInfo(std::string_view path) const;
+
    /// Prints the internal structure of this RFile to the given stream.
    void Print(std::ostream &out = std::cout) const;
 };

io/io/src/RFile.cxx

@@ -178,6 +178,17 @@ static void EnsureFileOpenAndBinary(const TFile *tfile, std::string_view path)
       throw ROOT::RException(R__FAIL("Opened file " + std::string(path) + " is not a ROOT binary file"));
 }
 
+static std::string ReconstructFullKeyPath(const TKey &key)
+{
+   std::string path = key.GetName();
+   TDirectory *parent = key.GetMotherDir();
+   while (parent && parent->GetMotherDir()) {
+      path = std::string(parent->GetName()) + "/" + path;
+      parent = parent->GetMotherDir();
+   }
+   return path;
+}
+
 /////////////////////////////////////////////////////////////////////////////////////////////////
 std::pair<std::string_view, std::string_view> ROOT::Experimental::Detail::DecomposePath(std::string_view path)
 {
@@ -269,21 +280,27 @@ TKey *RFile::GetTKey(std::string_view path) const
    return key;
 }
 
-void *RFile::GetUntyped(std::string_view pathSV, const std::type_info &type) const
+void *RFile::GetUntyped(std::string_view path,
+                        std::variant<const char *, std::reference_wrapper<const std::type_info>> type) const
 {
    if (!fFile)
       throw ROOT::RException(R__FAIL("File has been closed"));
 
-   std::string path{pathSV};
+   std::string pathStr{path};
+
+   struct {
+      TClass *operator()(const char *name) { return TClass::GetClass(name); }
+      TClass *operator()(std::reference_wrapper<const std::type_info> ty) { return TClass::GetClass(ty.get()); }
+   } typeVisitor;
+   const TClass *cls = std::visit(std::move(typeVisitor), type);
 
-   const TClass *cls = TClass::GetClass(type);
    if (!cls)
-      throw ROOT::RException(R__FAIL(std::string("Could not determine type of object ") + path));
+      throw ROOT::RException(R__FAIL(std::string("Could not determine type of object ") + pathStr));
 
-   if (auto err = ValidateAndNormalizePath(path); !err.empty())
-      throw RException(R__FAIL("Invalid object path '" + path + "': " + err));
+   if (auto err = ValidateAndNormalizePath(pathStr); !err.empty())
+      throw RException(R__FAIL("Invalid object pathStr '" + pathStr + "': " + err));
 
-   TKey *key = GetTKey(path);
+   TKey *key = GetTKey(pathStr);
    void *obj = key ? key->ReadObjectAny(cls) : nullptr;
 
    if (obj) {
@@ -513,3 +530,24 @@ void RFile::Close()
    // NOTE: this also flushes the file internally
    fFile.reset();
 }
+
+std::optional<ROOT::Experimental::RKeyInfo> RFile::GetKeyInfo(std::string_view path) const
+{
+   const TKey *key = GetTKey(path);
+   if (!key)
+      return {};
+
+   RKeyInfo keyInfo;
+   keyInfo.fPath = ReconstructFullKeyPath(*key);
+   keyInfo.fClassName = key->GetClassName();
+   keyInfo.fCycle = key->GetCycle();
+   keyInfo.fTitle = key->GetTitle();
+
+   return keyInfo;
+}
+
+void *ROOT::Experimental::Internal::RFile_GetObjectFromKey(RFile &file, const RKeyInfo &key)
+{
+   void *obj = file.GetUntyped(key.GetPath(), key.GetClassName().c_str());
+   return obj;
+}

io/io/test/CMakeLists.txt

@@ -19,6 +19,9 @@ if(uring AND NOT DEFINED ENV{ROOTTEST_IGNORE_URING})
 endif()
 
 ROOT_ADD_GTEST(rfile rfile.cxx LIBRARIES RIO Hist)
+if(pyroot)
+  ROOT_ADD_PYUNITTEST(rfile_py rfile.py)
+endif()
 
 # Temporarily disabled.  Test routinely fails on MacOS and some Linuxes.
 #if(NOT WIN32 AND (NOT MACOS_VERSION OR NOT MACOSX_VERSION VERSION_LESS 13.00))

io/io/test/rfile.cxx

@@ -701,3 +701,26 @@ TEST(RFile, NormalizedPaths)
    EXPECT_NE(file->Get<TH1D>("//a////b/c"), nullptr);
    EXPECT_THROW(file->Get<TH1D>("a/b/c/"), ROOT::RException);
 }
+
+TEST(RFile, GetKeyInfo)
+{
+   FileRaii fileGuard("test_rfile_getkeyinfo.root");
+
+   auto file = RFile::Recreate(fileGuard.GetPath());
+   std::string obj = "obj";
+   file->Put("/s", obj);
+   file->Put("a/b/c", obj);
+   file->Put("b", obj);
+   file->Put("/a/d", obj);
+
+   EXPECT_EQ(file->GetKeyInfo("foo"), std::nullopt);
+
+   for (const std::string_view path : { "/s", "a/b/c", "b", "/a/d" }) {
+      auto key = file->GetKeyInfo(path);
+      ASSERT_NE(key, std::nullopt);
+      EXPECT_EQ(key->GetPath(), path[0] == '/' ? path.substr(1) : path);
+      EXPECT_EQ(key->GetClassName(), "string");
+      EXPECT_EQ(key->GetTitle(), "");
+      EXPECT_EQ(key->GetCycle(), 1);
+   }
+}