[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

Author: Grigori Rybkine

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__

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)
+            for i in range(nentries):
+                entry["f"] = i
+                entry["mystr"] = f"{i} string stored in RNTuple"
+                writer.Fill(entry)
         # The model should not have been destroyed (a clone has been used).
         self.assertFalse(model.IsFrozen())
 
-        # Accessing the writer after the context manager is an error
-        with self.assertRaisesRegex(RuntimeError, "cannot access RNTupleWriter after"):
+        # Upon exiting the context, the writer is destructed
+        with self.assertRaisesRegex(ReferenceError, "attempt to access a null-pointer"):
             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")
+
+        # Upon exiting the context, the reader is destructed
+        with self.assertRaisesRegex(ReferenceError, "attempt to access a null-pointer"):
             reader.GetNEntries()
 
+        # Last entry values are still accessible after the reader is destructed
+        self.assertEqual(entry["f"], nentries - 1)
+        self.assertEqual(entry["mystr"], f"{nentries - 1} string stored in RNTuple")
+
     def test_write_fields(self):
         """Can create writer with on-the-fly model"""
 
@@ -73,14 +77,15 @@ def test_append_open(self):
         self.assertFalse(model.IsFrozen())
 
         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,8 +121,12 @@ 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"
@@ -127,49 +135,21 @@ def test_nested_ctxmanager(self):
             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)
-
-            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"""
-
-        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)
+                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)
 
             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 reader as r2:
+                    print(r2.GetNEntries())
+                with self.assertRaisesRegex(ReferenceError, "attempt to access a null-pointer"):
+                    print(r1.GetNEntries())
 
         finally:
             import os
             os.remove(fileName)
-                    
-