From 830435f3459a774dfdfbd385e1bd22ad75454dad Mon Sep 17 00:00:00 2001 From: Jonas Hahnfeld Date: Wed, 27 Nov 2024 15:52:41 +0100 Subject: [PATCH] [ntuple] Pythonize context manager for RNTupleWriter --- .../python/ROOT/_pythonization/_rntuple.py | 8 ++++++++ tree/ntuple/v7/test/ntuple_basics.py | 9 ++++----- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py index e5d172daba9fa..e17d4ff06b01f 100644 --- 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 index 86005fd12a1c6..4890a5454042c 100644 --- 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)