Skip to content

Commit a460d3f

Browse files
silverweeddpiparo
silverweed
authored and
dpiparo
committed
[rfile] add first pythonizations
1 parent e6a96b0 commit a460d3f

File tree

6 files changed

+307
-8
lines changed

6 files changed

+307
-8
lines changed

‎bindings/pyroot/pythonizations/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ set(py_sources
8686
ROOT/_pythonization/_generic.py
8787
ROOT/_pythonization/_memory_utils.py
8888
ROOT/_pythonization/_pyz_utils.py
89+
ROOT/_pythonization/_rfile.py
8990
ROOT/_pythonization/_rntuple.py
9091
ROOT/_pythonization/_runtime_error.py
9192
ROOT/_pythonization/_rvec.py

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

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
# Author: Giacomo Parolini CERN 04/2025
2+
3+
r"""
4+
\pythondoc RFile
5+
6+
TODO: document RFile
7+
8+
\code{.py}
9+
# TODO code example
10+
\endcode
11+
12+
\endpythondoc
13+
"""
14+
15+
from . import pythonization
16+
17+
18+
class _RFile_Get:
19+
"""
20+
Allow access to objects through the method Get().
21+
This is pythonized to allow Get() to be called both with and without a template argument.
22+
"""
23+
24+
def __init__(self, rfile):
25+
self._rfile = rfile
26+
27+
def __call__(self, namecycle):
28+
"""
29+
Non-templated Get()
30+
"""
31+
import ROOT
32+
import cppyy
33+
34+
key = self._rfile.GetKeyInfo(namecycle)
35+
if key:
36+
obj = ROOT.Experimental.Internal.RFile_GetObjectFromKey(self._rfile, key)
37+
return cppyy.bind_object(obj, key.GetClassName())
38+
# No key
39+
return None
40+
41+
def __getitem__(self, template_arg):
42+
"""
43+
Templated Get()
44+
"""
45+
46+
def getitem_wrapper(namecycle):
47+
obj = self._rfile._OriginalGet[template_arg](namecycle)
48+
return obj if obj else None
49+
50+
return getitem_wrapper
51+
52+
53+
class _RFile_Put:
54+
"""
55+
Allow writing objects through the method Put().
56+
This is pythonized to allow Put() to be called both with and without a template argument.
57+
"""
58+
59+
def __init__(self, rfile):
60+
self._rfile = rfile
61+
62+
def __call__(self, name, obj):
63+
"""
64+
Non-templated Put()
65+
"""
66+
objType = type(obj)
67+
if objType == str:
68+
# special case: automatically convert python str to std::string
69+
className = "std::string"
70+
elif not hasattr(objType, '__cpp_name__'):
71+
raise TypeError(f"type {objType} is not supported by ROOT I/O")
72+
else:
73+
className = objType.__cpp_name__
74+
self._rfile.Put[className](name, obj)
75+
76+
def __getitem__(self, template_arg):
77+
"""
78+
Templated Put()
79+
"""
80+
return self._rfile._OriginalPut[template_arg]
81+
82+
83+
def _RFileExit(obj, exc_type, exc_val, exc_tb):
84+
"""
85+
Close the RFile object.
86+
Signature and return value are imposed by Python, see
87+
https://docs.python.org/3/library/stdtypes.html#typecontextmanager.
88+
"""
89+
obj.Close()
90+
return False
91+
92+
93+
def _RFileOpen(original):
94+
"""
95+
Pythonization for the factory methods (Recreate, Open, Update)
96+
"""
97+
98+
def rfile_open_wrapper(klass, *args):
99+
rfile = original(*args)
100+
rfile._OriginalGet = rfile.Get
101+
rfile.Get = _RFile_Get(rfile)
102+
rfile._OriginalPut = rfile.Put
103+
rfile.Put = _RFile_Put(rfile)
104+
return rfile
105+
106+
return rfile_open_wrapper
107+
108+
109+
def _RFileInit(rfile):
110+
"""
111+
Prevent the creation of RFile through constructor (must use a factory method)
112+
"""
113+
raise NotImplementedError("RFile can only be created via Recreate, Open or Update")
114+
115+
116+
def _GetKeyInfo(rfile, path):
117+
key = rfile._OriginalGetKeyInfo(path)
118+
if key.has_value():
119+
return key.value()
120+
return None
121+
122+
123+
def _ListKeys(rfile, basePath="", listObjects = True, listDirs = False, listRecursive = True):
124+
from ROOT.Experimental import RFile
125+
126+
flags = (listObjects * RFile.kListObjects) | (listDirs * RFile.kListDirs) | (listRecursive * RFile.kListRecursive)
127+
iter = rfile._OriginalListKeys(basePath, flags)
128+
return iter
129+
130+
131+
@pythonization("RFile", ns="ROOT::Experimental")
132+
def pythonize_rfile(klass):
133+
# Explicitly prevent to create a RFile via ctor
134+
klass.__init__ = _RFileInit
135+
136+
# Pythonize factory methods
137+
klass.Open = classmethod(_RFileOpen(klass.Open))
138+
klass.Update = classmethod(_RFileOpen(klass.Update))
139+
klass.Recreate = classmethod(_RFileOpen(klass.Recreate))
140+
141+
# Pythonization for __enter__ and __exit__ methods
142+
# These make RFile usable in a `with` statement as a context manager
143+
klass.__enter__ = lambda rfile: rfile
144+
klass.__exit__ = _RFileExit
145+
klass._OriginalGetKeyInfo = klass.GetKeyInfo
146+
klass.GetKeyInfo = _GetKeyInfo
147+
klass._OriginalListKeys = klass.ListKeys
148+
klass.ListKeys = _ListKeys

‎io/io/inc/ROOT/RFile.hxx‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
#include <memory>
1717
#include <string_view>
1818
#include <typeinfo>
19+
#include <variant>
1920

2021
class TFile;
2122
class TIterator;
@@ -24,6 +25,7 @@ class TKey;
2425
namespace ROOT {
2526
namespace Experimental {
2627

28+
class RKeyInfo;
2729
class RFile;
2830

2931
namespace Internal {
@@ -238,7 +240,8 @@ class RFile final {
238240

239241
/// Gets object `path` from the file and returns an **owning** pointer to it.
240242
/// The caller should immediately wrap it into a unique_ptr of the type described by `type`.
241-
[[nodiscard]] void *GetUntyped(std::string_view path, const std::type_info &type) const;
243+
[[nodiscard]] void *GetUntyped(std::string_view path,
244+
std::variant<const char *, std::reference_wrapper<const std::type_info>> type) const;
242245

243246
/// Writes `obj` to file, without taking its ownership.
244247
void PutUntyped(std::string_view path, const std::type_info &type, const void *obj, std::uint32_t flags);

‎io/io/src/RFile.cxx‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -280,21 +280,27 @@ TKey *RFile::GetTKey(std::string_view path) const
280280
return key;
281281
}
282282

283-
void *RFile::GetUntyped(std::string_view pathSV, const std::type_info &type) const
283+
void *RFile::GetUntyped(std::string_view path,
284+
std::variant<const char *, std::reference_wrapper<const std::type_info>> type) const
284285
{
285286
if (!fFile)
286287
throw ROOT::RException(R__FAIL("File has been closed"));
287288

288-
std::string path{pathSV};
289+
std::string pathStr{path};
290+
291+
struct {
292+
TClass *operator()(const char *name) { return TClass::GetClass(name); }
293+
TClass *operator()(std::reference_wrapper<const std::type_info> ty) { return TClass::GetClass(ty.get()); }
294+
} typeVisitor;
295+
const TClass *cls = std::visit(std::move(typeVisitor), type);
289296

290-
const TClass *cls = TClass::GetClass(type);
291297
if (!cls)
292-
throw ROOT::RException(R__FAIL(std::string("Could not determine type of object ") + path));
298+
throw ROOT::RException(R__FAIL(std::string("Could not determine type of object ") + pathStr));
293299

294-
if (auto err = ValidateAndNormalizePath(path); !err.empty())
295-
throw RException(R__FAIL("Invalid object path '" + path + "': " + err));
300+
if (auto err = ValidateAndNormalizePath(pathStr); !err.empty())
301+
throw RException(R__FAIL("Invalid object pathStr '" + pathStr + "': " + err));
296302

297-
TKey *key = GetTKey(path);
303+
TKey *key = GetTKey(pathStr);
298304
void *obj = key ? key->ReadObjectAny(cls) : nullptr;
299305

300306
if (obj) {

‎io/io/test/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@ if(uring AND NOT DEFINED ENV{ROOTTEST_IGNORE_URING})
1919
endif()
2020

2121
ROOT_ADD_GTEST(rfile rfile.cxx LIBRARIES RIO Hist)
22+
if(pyroot)
23+
ROOT_ADD_PYUNITTEST(rfile_py rfile.py)
24+
endif()
2225

2326
# Temporarily disabled. Test routinely fails on MacOS and some Linuxes.
2427
#if(NOT WIN32 AND (NOT MACOS_VERSION OR NOT MACOSX_VERSION VERSION_LESS 13.00))

‎io/io/test/rfile.py‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
import os
2+
import unittest
3+
import ROOT
4+
5+
RFile = ROOT.Experimental.RFile
6+
7+
8+
class RFileTests(unittest.TestCase):
9+
def test_open_for_reading(self):
10+
"""A RFile can read a ROOT file created by TFile"""
11+
12+
fileName = "test_rfile_read_py.root"
13+
14+
# Create a root file to open
15+
with ROOT.TFile.Open(fileName, "RECREATE") as tfile:
16+
hist = ROOT.TH1D("hist", "", 100, -10, 10)
17+
hist.FillRandom("gaus", 100)
18+
tfile.WriteObject(hist, "hist")
19+
20+
with RFile.Open(fileName) as rfile:
21+
hist = rfile.Get("hist")
22+
self.assertNotEqual(hist, None)
23+
self.assertEqual(rfile.Get[ROOT.TH1D]("inexistent"), None)
24+
self.assertEqual(rfile.Get[ROOT.TH1F]("hist"), None)
25+
self.assertNotEqual(rfile.Get[ROOT.TH1]("hist"), None)
26+
27+
with self.assertRaises(ROOT.RException):
28+
# This should fail because the file was opened as read-only
29+
rfile.Put("foo", hist)
30+
31+
os.remove(fileName)
32+
33+
def test_writing_reading(self):
34+
"""A RFile can be written into and read from"""
35+
36+
fileName = "test_rfile_writeread_py.root"
37+
38+
with RFile.Recreate(fileName) as rfile:
39+
hist = ROOT.TH1D("hist", "", 100, -10, 10)
40+
hist.FillRandom("gaus", 10)
41+
rfile.Put("hist", hist)
42+
with self.assertRaises(ROOT.RException):
43+
rfile.Put("hist/2", hist)
44+
45+
with RFile.Open(fileName) as rfile:
46+
hist = rfile.Get("hist")
47+
self.assertNotEqual(hist, None)
48+
49+
os.remove(fileName)
50+
51+
def test_getkeyinfo(self):
52+
"""A RFile can query individual keys of its objects"""
53+
54+
fileName = "test_rfile_getkeyinfo_py.root"
55+
56+
with RFile.Recreate(fileName) as rfile:
57+
hist = ROOT.TH1D("hist", "", 100, -10, 10)
58+
hist.FillRandom("gaus", 10)
59+
rfile.Put("hist", hist)
60+
rfile.Put("foo/hist", hist)
61+
rfile.Put("foo/bar/hist", hist)
62+
rfile.Put("foo/bar/hist2", hist)
63+
rfile.Put("foo/hist2", hist)
64+
65+
with RFile.Open(fileName) as rfile:
66+
key = rfile.GetKeyInfo("hist")
67+
self.assertEqual(key.GetPath(), "hist")
68+
self.assertEqual(key.GetClassName(), "TH1D")
69+
70+
key = rfile.GetKeyInfo("does_not_exist")
71+
self.assertEqual(key, None)
72+
73+
def test_listkeys(self):
74+
"""A RFile can query the keys of its objects and directories"""
75+
76+
fileName = "test_rfile_listkeys_py.root"
77+
78+
with RFile.Recreate(fileName) as rfile:
79+
hist = ROOT.TH1D("hist", "", 100, -10, 10)
80+
hist.FillRandom("gaus", 10)
81+
rfile.Put("hist", hist)
82+
rfile.Put("foo/hist", hist)
83+
rfile.Put("foo/bar/hist", hist)
84+
rfile.Put("foo/bar/hist2", hist)
85+
rfile.Put("foo/hist2", hist)
86+
87+
with RFile.Open(fileName) as rfile:
88+
keys = [key.GetPath() for key in rfile.ListKeys()]
89+
self.assertEqual(keys, ["hist", "foo/hist", "foo/bar/hist", "foo/bar/hist2", "foo/hist2"])
90+
91+
keys = [key.GetClassName() for key in rfile.ListKeys()]
92+
self.assertEqual(keys, ["TH1D"] * len(keys))
93+
94+
self.assertEqual(
95+
[key.GetPath() for key in rfile.ListKeys("foo")],
96+
["foo/hist", "foo/bar/hist", "foo/bar/hist2", "foo/hist2"],
97+
)
98+
99+
self.assertEqual([key.GetPath() for key in rfile.ListKeys("foo/bar")], ["foo/bar/hist", "foo/bar/hist2"])
100+
101+
self.assertEqual(
102+
[key.GetPath() for key in rfile.ListKeys("", listDirs=True, listObjects=False)], ["foo", "foo/bar"]
103+
)
104+
self.assertEqual(
105+
[key.GetPath() for key in rfile.ListKeys("", listDirs=True, listObjects=False, listRecursive=False)],
106+
["foo"],
107+
)
108+
self.assertEqual(
109+
[key.GetPath() for key in rfile.ListKeys("", listDirs=True, listRecursive=False)], ["hist", "foo"]
110+
)
111+
112+
os.remove(fileName)
113+
114+
def test_putUnsupportedType(self):
115+
fileName = "test_rfile_putunsupported_py.root"
116+
117+
with RFile.Recreate(fileName) as rfile:
118+
# Storing integers is unsupported
119+
with self.assertRaises(TypeError):
120+
rfile.Put("foo", 2)
121+
122+
# Storing lists without an explicit template is unsupported
123+
with self.assertRaises(TypeError):
124+
rfile.Put("bar", [2, 3])
125+
126+
# Storing lists with an explicit template is supported
127+
rfile.Put["std::vector<int>"]("bar", [2, 3])
128+
129+
# Storing strings is supported
130+
rfile.Put("str", "foobar")
131+
132+
with RFile.Open(fileName) as rfile:
133+
self.assertEqual(rfile.Get("str"), b"foobar")
134+
135+
os.remove(fileName)
136+
137+
if __name__ == "__main__":
138+
unittest.main()

0 commit comments

Comments
 (0)