From 52bc377c0b104284dba9aed99595cf0b237a9fdf Mon Sep 17 00:00:00 2001 From: Grigori Rybkine Date: Fri, 29 May 2026 19:22:40 +0200 Subject: [PATCH] [ntuple][python][ATLAS experiment] Re-Implement context management protocol for RNTupleReader/Writer bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py: add __enter__ method - returns self (an instance of RNTupleReader/RNTupleWriter), __exit__ method - calls RNTupleReader/RNTupleWriter destructor (if not destructed yet). tree/ntuple/test/ntuple_basics.py: update tests --- .../python/ROOT/_pythonization/_rntuple.py | 73 +++------- tree/ntuple/test/ntuple_basics.py | 130 ++++++++---------- 2 files changed, 81 insertions(+), 122 deletions(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py index 446476971151e..3b0b0a84726b9 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py @@ -124,51 +124,6 @@ def pythonize_RNTupleModel(klass): klass.GetDefaultEntry = _RNTupleModel_GetDefaultEntry klass.MakeField = MethodTemplateGetter(klass.MakeField, _RNTupleModel_MakeField) - - -# Wrapper class used for RNTupleReader and RNTupleWriter. -# It deletes the underlying smart pointer on context manager exit and ensures that the inner object becomes -# inaccessible by raising an error every time an attribute of the object is accessed. -# It also raises an error if `with` statements using the same object are nested. -# This is a generic class and can in principle be used with any class that needs this behavior. -class RNTupleContextWrapper: - def __init__(self, inner, pretty_name, on_ctx_enter = None, on_ctx_exit = None): - self._inner = inner - self._pretty_name = pretty_name - self._closed = False - self._in_context = False - self._on_ctx_enter = on_ctx_enter - self._on_ctx_exit = on_ctx_exit - - def __getattribute__(self, name): - if name.startswith('_'): - return super().__getattribute__(name) - - if super().__getattribute__("_closed"): - raise RuntimeError( - f"cannot access {super().__getattribute__('_pretty_name')} after the `with` statement is exited" - ) - return super().__getattribute__("_inner").__getattribute__(name) - - def __enter__(self, *args): - if self._on_ctx_enter: - self._on_ctx_enter(self._inner) - if self._closed: - raise RuntimeError(f"cannot reuse {self._pretty_name} in multiple `with` statements") - if self._in_context: - raise RuntimeError(f"cannot nest `with` statements using the same {self._pretty_name}") - - self._in_context = True - return self - - def __exit__(self, *args): - assert self._in_context and not self._closed - if self._on_ctx_exit: - self._on_ctx_exit(self._inner) - self._in_context = False - self._closed = True - self._inner.__smartptr__().reset() - return False def _RNTupleReader_Open(maybe_model, *args): @@ -177,7 +132,7 @@ def _RNTupleReader_Open(maybe_model, *args): maybe_model = maybe_model.Clone() import ROOT - return RNTupleContextWrapper(ROOT.RNTupleReader._Open(maybe_model, *args), "RNTupleReader") + return ROOT.RNTupleReader._Open(maybe_model, *args) def _RNTupleReader_LoadEntry(self, *args): @@ -186,6 +141,18 @@ def _RNTupleReader_LoadEntry(self, *args): return self._LoadEntry(*args) +def _RNTupleReaderWriter___enter__(self): + """Context management protocol. Returns self (an instance of RNTupleReader/RNTupleWriter).""" + if not self.__smartptr__(): + raise ValueError(f"I/O operation on destructed {type(self).__name__!r}.") + return self + + +def _RNTupleReaderWriter___exit__(self, *args): + """Context management protocol. Calls RNTupleReader/RNTupleWriter destructor (if not destructed yet).""" + self.__smartptr__().reset() + + @pythonization("RNTupleReader", ns="ROOT") def pythonize_RNTupleReader(klass): klass._Open = klass.Open @@ -194,13 +161,16 @@ def pythonize_RNTupleReader(klass): klass._LoadEntry = klass.LoadEntry klass.LoadEntry = _RNTupleReader_LoadEntry + klass.__enter__ = _RNTupleReaderWriter___enter__ + klass.__exit__ = _RNTupleReaderWriter___exit__ + def _RNTupleWriter_Append(model, *args): # In Python, the user cannot create REntries directly from a model, so we can safely clone it and avoid destructively passing the user argument. model = model.Clone() import ROOT - return RNTupleContextWrapper(ROOT.RNTupleWriter._Append(model, *args), "RNTupleWriter", on_ctx_exit = _RNTupleWriter_exit) + return ROOT.RNTupleWriter._Append(model, *args) def _RNTupleWriter_Recreate(model_or_fields, *args): @@ -209,7 +179,7 @@ def _RNTupleWriter_Recreate(model_or_fields, *args): model_or_fields = model_or_fields.Clone() import ROOT - return RNTupleContextWrapper(ROOT.RNTupleWriter._Recreate(model_or_fields, *args), "RNTupleWriter", on_ctx_exit = _RNTupleWriter_exit) + return ROOT.RNTupleWriter._Recreate(model_or_fields, *args) def _RNTupleWriter_Fill(self, *args): @@ -218,10 +188,6 @@ def _RNTupleWriter_Fill(self, *args): return self._Fill(*args) -def _RNTupleWriter_exit(self): - self.CommitDataset() - - @pythonization("RNTupleWriter", ns="ROOT") def pythonize_RNTupleWriter(klass): klass._Append = klass.Append @@ -231,3 +197,6 @@ def pythonize_RNTupleWriter(klass): klass._Fill = klass.Fill klass.Fill = _RNTupleWriter_Fill + + klass.__enter__ = _RNTupleReaderWriter___enter__ + klass.__exit__ = _RNTupleReaderWriter___exit__ diff --git a/tree/ntuple/test/ntuple_basics.py b/tree/ntuple/test/ntuple_basics.py index 4a1c0ec30b3d6..42b345a450032 100644 --- a/tree/ntuple/test/ntuple_basics.py +++ b/tree/ntuple/test/ntuple_basics.py @@ -14,33 +14,37 @@ def test_write_read(self): model.MakeField["int"]("f") model.MakeField["std::string"]("mystr") + nentries = 2 with ROOT.RNTupleWriter.Recreate(model, "ntpl", "test_ntuple_py_write_read.root") as writer: entry = writer.CreateEntry() - entry["f"] = 42 - entry["mystr"] = "string stored in RNTuple" - writer.Fill(entry) - # The model should not have been destroyed (a clone has been used). - self.assertFalse(model.IsFrozen()) + for i in range(nentries): + entry["f"] = i + entry["mystr"] = f"{i} string stored in RNTuple" + writer.Fill(entry) + self.assertFalse(model.IsFrozen(), + msg="The model should not have been destroyed (a clone has been used).") - # Accessing the writer after the context manager is an error - with self.assertRaisesRegex(RuntimeError, "cannot access RNTupleWriter after"): + with self.assertRaisesRegex(ReferenceError, "attempt to access a null-pointer", + msg="Upon exiting the context, the writer is destructed."): writer.GetNEntries() with ROOT.RNTupleReader.Open("ntpl", "test_ntuple_py_write_read.root") as reader: - self.assertEqual(reader.GetNEntries(), 1) + self.assertEqual(reader.GetNEntries(), nentries) entry = reader.CreateEntry() - reader.LoadEntry(0, entry) - self.assertEqual(entry["f"], 42) - self.assertEqual(entry["mystr"], "string stored in RNTuple") - - # Entry values are still accessible after the reader is gone - self.assertEqual(entry["f"], 42) - self.assertEqual(entry["mystr"], "string stored in RNTuple") - - # Accessing the reader after the context manager is an error - with self.assertRaisesRegex(RuntimeError, "cannot access RNTupleReader after"): + for i in reader: + reader.LoadEntry(i, entry) + with self.subTest(i=i): + self.assertEqual(entry["f"], i) + self.assertEqual(entry["mystr"], f"{i} string stored in RNTuple") + + with self.assertRaisesRegex(ReferenceError, "attempt to access a null-pointer", + msg="Upon exiting the context, the reader is destructed."): reader.GetNEntries() + msg = "Last entry values are still accessible after the reader is destructed." + self.assertEqual(entry["f"], nentries - 1, msg=msg) + self.assertEqual(entry["mystr"], f"{nentries - 1} string stored in RNTuple", msg=msg) + def test_write_fields(self): """Can create writer with on-the-fly model""" @@ -69,18 +73,19 @@ def test_append_open(self): entry["f"] = 42 writer.Fill(entry) - # The model should not have been destroyed (a clone has been used). - self.assertFalse(model.IsFrozen()) + self.assertFalse(model.IsFrozen(), + msg="The model should not have been destroyed (a clone has been used).") with ROOT.TFile.Open("test_ntuple_py_append.root") as f: - reader = ROOT.RNTupleReader.Open(f["ntpl"]) - self.assertEqual(reader.GetNEntries(), 1) - entry = reader.CreateEntry() - reader.LoadEntry(0, entry) - self.assertEqual(entry["f"], 42) + with ROOT.RNTupleReader.Open(f["ntpl"]) as reader: + self.assertEqual(reader.GetNEntries(), 1) + entry = reader.CreateEntry() + reader.LoadEntry(0, entry) + self.assertEqual(entry["f"], 42) - # Entry values are still accessible after the reader is gone - self.assertEqual(entry["f"], 42) + with self.subTest(repr(reader)): + self.assertFalse(reader, "RNTupleReader destructed") + self.assertEqual(entry["f"], 42, "Entry values still accessible") def test_read_model(self): """Can impose a model when reading.""" @@ -100,8 +105,7 @@ def test_read_model(self): entry = reader.CreateEntry() if not platform.system() == "Windows": # TODO: re-enable it on Windows once the exception handling is fixed - with self.assertRaises(Exception): - # Field f2 does not exist in imposed model + with self.assertRaises(ROOT.RException, msg="Field f2 does not exist in imposed model"): entry["f2"] = 42 def test_forbid_writing_wrong_type(self): @@ -117,59 +121,45 @@ class WrongClass: ... with self.assertRaises(TypeError): entry["mystr"] = WrongClass() - def test_nested_ctxmanager(self): - """Nesting context managers of the same object is an error""" + def test_singleuse_ctxmanager(self): + """RNTupleReader/RNTupleWriter context managers are single use context managers. + + Upon exiting the context, they are destructed. + They are not reentrant - cannot be used in nested 'with' statements, + or are not reusable - cannot be used multiple times.""" try: - fileName = "test_ntuple_nested_ctxmanager_py.root" + fileName = "test_singleuse_ctxmanager_py.root" model = ROOT.RNTupleModel.Create() model.MakeField["int"]("f") writer = ROOT.RNTupleWriter.Recreate(model, "ntpl", fileName) with writer as w1: entry1 = w1.CreateEntry() - with self.assertRaisesRegex(RuntimeError, "cannot nest `with`"): - with writer as w2: - entry2 = w2.CreateEntry() - entry1["f"] = 2 - entry2["f"] = 4 - w2.Fill(entry2) - w1.Fill(entry1) + entry1["f"] = 2 + with writer as w2: + entry2 = w2.CreateEntry() + entry2["f"] = 4 + w2.Fill(entry2) + with self.assertRaisesRegex((ReferenceError, TypeError), "attempt to access a null-pointer"): + w1.Fill(entry1) + + with self.assertRaisesRegex(ValueError, "I/O operation on destructed 'RNTupleWriter'"): + with writer as w: + entry = w.CreateEntry() + entry["f"] = 8 + w.Fill(entry) reader = ROOT.RNTupleReader.Open("ntpl", fileName) with reader as r1: - with self.assertRaisesRegex(RuntimeError, "cannot nest `with`"): - with reader as r2: - print(r2.GetNEntries()) - print(r1.GetNEntries()) - - finally: - import os - os.remove(fileName) - - def test_weird_ctxmanager(self): - """Using an existing object with a context manager""" + with reader as r2: + print(r2.GetNEntries()) + with self.assertRaisesRegex(ReferenceError, "attempt to access a null-pointer"): + print(r1.GetNEntries()) - try: - fileName = "test_ntuple_weird_ctxmanager_py.root" - model = ROOT.RNTupleModel.Create() - model.MakeField["int"]("f") - writer = ROOT.RNTupleWriter.Recreate(model, "ntpl", fileName) - entry = writer.CreateEntry() - with writer as w1: - w1.Fill(entry) - - with self.assertRaisesRegex(RuntimeError, "after the `with` statement"): - writer.Fill(entry) - - reader = ROOT.RNTupleReader.Open("ntpl", fileName) - with reader as r1: - with self.assertRaisesRegex(RuntimeError, "cannot nest `with`"): - with reader as r2: - print(r2.GetNEntries()) - print(r1.GetNEntries()) + with self.assertRaisesRegex(ValueError, "I/O operation on destructed 'RNTupleReader'"): + with reader as r: + print(r.GetNEntries()) finally: import os os.remove(fileName) - -