diff --git a/bindings/pyroot/pythonizations/CMakeLists.txt b/bindings/pyroot/pythonizations/CMakeLists.txt index af2f6a39119b4..5af09f8a1a335 100644 --- a/bindings/pyroot/pythonizations/CMakeLists.txt +++ b/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 diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py new file mode 100644 index 0000000000000..d29d713966a73 --- /dev/null +++ b/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 diff --git a/io/io/inc/ROOT/RFile.hxx b/io/io/inc/ROOT/RFile.hxx index 6188961e95f3c..7f153f74b1b38 100644 --- a/io/io/inc/ROOT/RFile.hxx +++ b/io/io/inc/ROOT/RFile.hxx @@ -11,10 +11,12 @@ #include #include -#include +#include #include +#include #include #include +#include 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("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> 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 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; }; diff --git a/io/io/src/RFile.cxx b/io/io/src/RFile.cxx index c5364fed3eb47..541f2569cc10f 100644 --- a/io/io/src/RFile.cxx +++ b/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 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> 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 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 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; +} diff --git a/io/io/test/CMakeLists.txt b/io/io/test/CMakeLists.txt index 06695660ef290..8e1031f01a1e9 100644 --- a/io/io/test/CMakeLists.txt +++ b/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)) diff --git a/io/io/test/rfile.cxx b/io/io/test/rfile.cxx index 526172a4964cf..2d40607e6e703 100644 --- a/io/io/test/rfile.cxx +++ b/io/io/test/rfile.cxx @@ -701,3 +701,26 @@ TEST(RFile, NormalizedPaths) EXPECT_NE(file->Get("//a////b/c"), nullptr); EXPECT_THROW(file->Get("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); + } +} diff --git a/io/io/test/rfile.py b/io/io/test/rfile.py new file mode 100644 index 0000000000000..4f23d5a0964f8 --- /dev/null +++ b/io/io/test/rfile.py @@ -0,0 +1,138 @@ +import os +import unittest +import ROOT + +RFile = ROOT.Experimental.RFile + + +class RFileTests(unittest.TestCase): + def test_open_for_reading(self): + """A RFile can read a ROOT file created by TFile""" + + fileName = "test_rfile_read_py.root" + + # Create a root file to open + with ROOT.TFile.Open(fileName, "RECREATE") as tfile: + hist = ROOT.TH1D("hist", "", 100, -10, 10) + hist.FillRandom("gaus", 100) + tfile.WriteObject(hist, "hist") + + with RFile.Open(fileName) as rfile: + hist = rfile.Get("hist") + self.assertNotEqual(hist, None) + self.assertEqual(rfile.Get[ROOT.TH1D]("inexistent"), None) + self.assertEqual(rfile.Get[ROOT.TH1F]("hist"), None) + self.assertNotEqual(rfile.Get[ROOT.TH1]("hist"), None) + + with self.assertRaises(ROOT.RException): + # This should fail because the file was opened as read-only + rfile.Put("foo", hist) + + os.remove(fileName) + + def test_writing_reading(self): + """A RFile can be written into and read from""" + + fileName = "test_rfile_writeread_py.root" + + with RFile.Recreate(fileName) as rfile: + hist = ROOT.TH1D("hist", "", 100, -10, 10) + hist.FillRandom("gaus", 10) + rfile.Put("hist", hist) + with self.assertRaises(ROOT.RException): + rfile.Put("hist/2", hist) + + with RFile.Open(fileName) as rfile: + hist = rfile.Get("hist") + self.assertNotEqual(hist, None) + + os.remove(fileName) + + def test_getkeyinfo(self): + """A RFile can query individual keys of its objects""" + + fileName = "test_rfile_getkeyinfo_py.root" + + with RFile.Recreate(fileName) as rfile: + hist = ROOT.TH1D("hist", "", 100, -10, 10) + hist.FillRandom("gaus", 10) + rfile.Put("hist", hist) + rfile.Put("foo/hist", hist) + rfile.Put("foo/bar/hist", hist) + rfile.Put("foo/bar/hist2", hist) + rfile.Put("foo/hist2", hist) + + with RFile.Open(fileName) as rfile: + key = rfile.GetKeyInfo("hist") + self.assertEqual(key.GetPath(), "hist") + self.assertEqual(key.GetClassName(), "TH1D") + + key = rfile.GetKeyInfo("does_not_exist") + self.assertEqual(key, None) + + def test_listkeys(self): + """A RFile can query the keys of its objects and directories""" + + fileName = "test_rfile_listkeys_py.root" + + with RFile.Recreate(fileName) as rfile: + hist = ROOT.TH1D("hist", "", 100, -10, 10) + hist.FillRandom("gaus", 10) + rfile.Put("hist", hist) + rfile.Put("foo/hist", hist) + rfile.Put("foo/bar/hist", hist) + rfile.Put("foo/bar/hist2", hist) + rfile.Put("foo/hist2", hist) + + with RFile.Open(fileName) as rfile: + keys = [key.GetPath() for key in rfile.ListKeys()] + self.assertEqual(keys, ["hist", "foo/hist", "foo/bar/hist", "foo/bar/hist2", "foo/hist2"]) + + keys = [key.GetClassName() for key in rfile.ListKeys()] + self.assertEqual(keys, ["TH1D"] * len(keys)) + + self.assertEqual( + [key.GetPath() for key in rfile.ListKeys("foo")], + ["foo/hist", "foo/bar/hist", "foo/bar/hist2", "foo/hist2"], + ) + + self.assertEqual([key.GetPath() for key in rfile.ListKeys("foo/bar")], ["foo/bar/hist", "foo/bar/hist2"]) + + self.assertEqual( + [key.GetPath() for key in rfile.ListKeys("", listDirs=True, listObjects=False)], ["foo", "foo/bar"] + ) + self.assertEqual( + [key.GetPath() for key in rfile.ListKeys("", listDirs=True, listObjects=False, listRecursive=False)], + ["foo"], + ) + self.assertEqual( + [key.GetPath() for key in rfile.ListKeys("", listDirs=True, listRecursive=False)], ["hist", "foo"] + ) + + os.remove(fileName) + + def test_putUnsupportedType(self): + fileName = "test_rfile_putunsupported_py.root" + + with RFile.Recreate(fileName) as rfile: + # Storing integers is unsupported + with self.assertRaises(TypeError): + rfile.Put("foo", 2) + + # Storing lists without an explicit template is unsupported + with self.assertRaises(TypeError): + rfile.Put("bar", [2, 3]) + + # Storing lists with an explicit template is supported + rfile.Put["std::vector"]("bar", [2, 3]) + + # Storing strings is supported + rfile.Put("str", "foobar") + + with RFile.Open(fileName) as rfile: + self.assertEqual(rfile.Get("str"), b"foobar") + + os.remove(fileName) + +if __name__ == "__main__": + unittest.main()