Skip to content

[ntuple][python][ATLAS experiment] Re-Implement context management pr… #22432

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rybkine wants to merge 1 commit into
base: master
Choose a base branch
from
+81 −122
Open

[ntuple][python][ATLAS experiment] Re-Implement context management pr… #22432

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 21 additions & 52 deletions bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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):
Expand All @@ -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
Expand All @@ -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):
Expand All @@ -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):
Expand All @@ -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
Expand All @@ -231,3 +197,6 @@ def pythonize_RNTupleWriter(klass):

klass._Fill = klass.Fill
klass.Fill = _RNTupleWriter_Fill

klass.__enter__ = _RNTupleReaderWriter___enter__
klass.__exit__ = _RNTupleReaderWriter___exit__
130 changes: 60 additions & 70 deletions tree/ntuple/test/ntuple_basics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"""

Expand Down Expand Up @@ -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."""
Expand All @@ -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):
Expand All @@ -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)