[ntuple] Pythonize context manager for RNTupleWriter

root-project · Nov 27, 2024 · 830435f · 830435f
1 parent 046da80
commit 830435f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py
@@ -95,7 +95,15 @@ def _RNTupleWriter_Fill(self, *args):
     return self._Fill(*args)
 
 
+def _RNTupleWriter_exit(self, *args):
+    self.CommitDataset()
+    return False
+
+
 @pythonization("RNTupleWriter", ns="ROOT::Experimental")
 def pythonize_RNTupleWriter(klass):
     klass._Fill = klass.Fill
     klass.Fill = _RNTupleWriter_Fill
+
+    klass.__enter__ = lambda writer: writer
+    klass.__exit__ = _RNTupleWriter_exit
diff --git a/tree/ntuple/v7/test/ntuple_basics.py b/tree/ntuple/v7/test/ntuple_basics.py
@@ -14,11 +14,10 @@ def test_write_read(self):
 
         model = RNTupleModel.Create()
         model.MakeField["int"]("f")
-        writer = RNTupleWriter.Recreate(model, "ntpl", "test_ntuple_py_write_read.root")
-        entry = writer.CreateEntry()
-        entry["f"] = 42
-        writer.Fill(entry)
-        del writer
+        with RNTupleWriter.Recreate(model, "ntpl", "test_ntuple_py_write_read.root") as writer:
+            entry = writer.CreateEntry()
+            entry["f"] = 42
+            writer.Fill(entry)
 
         reader = RNTupleReader.Open("ntpl", "test_ntuple_py_write_read.root")
         self.assertEqual(reader.GetNEntries(), 1)