From 046da800ce3ad0828620df501ee702057ce415dd Mon Sep 17 00:00:00 2001
From: Jonas Hahnfeld <jonas.hahnfeld@cern.ch>
Date: Tue, 26 Nov 2024 13:28:16 +0100
Subject: [PATCH] [ntuple] Pythonize REntry

Make it behave like a dictionary that automatically gets the field type
and uses it internally.
---
 .../python/ROOT/_pythonization/_rntuple.py    | 30 +++++++++++++++++++
 tree/ntuple/v7/test/ntuple_basics.py          |  4 +++
 2 files changed, 34 insertions(+)

diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py
index 6902a7123d4eb..e5d172daba9fa 100644
--- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py
+++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rntuple.py
@@ -12,6 +12,36 @@
 from ._pyz_utils import MethodTemplateGetter, MethodTemplateWrapper
 
 
+def _REntry_GetPtr(self, key):
+    # key can be either a RFieldToken already or a string. In the latter case, get a token to use it twice.
+    if (
+        not hasattr(type(key), "__cpp_name__")
+        or type(key).__cpp_name__ != "ROOT::Experimental::REntry::RFieldToken"
+    ):
+        key = self.GetToken(key)
+    fieldType = self.GetTypeName(key)
+    return self._GetPtr[fieldType](key)
+
+
+def _REntry_getitem(self, key):
+    ptr = self.GetPtr(key)
+    return ptr.get()[0]
+
+
+def _REntry_setitem(self, key, value):
+    ptr = self.GetPtr(key)
+    ptr.get()[0] = value
+
+
+@pythonization("REntry", ns="ROOT::Experimental")
+def pythonize_REntry(klass):
+    klass._GetPtr = klass.GetPtr
+    klass.GetPtr = _REntry_GetPtr
+
+    klass.__getitem__ = _REntry_getitem
+    klass.__setitem__ = _REntry_setitem
+
+
 def _RNTupleModel_CreateBare(*args):
     if len(args) >= 1:
         raise ValueError("no support for passing explicit RFieldZero")
diff --git a/tree/ntuple/v7/test/ntuple_basics.py b/tree/ntuple/v7/test/ntuple_basics.py
index 6e93842d1dba5..86005fd12a1c6 100644
--- a/tree/ntuple/v7/test/ntuple_basics.py
+++ b/tree/ntuple/v7/test/ntuple_basics.py
@@ -16,8 +16,12 @@ def test_write_read(self):
         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
 
         reader = RNTupleReader.Open("ntpl", "test_ntuple_py_write_read.root")
         self.assertEqual(reader.GetNEntries(), 1)
+        entry = reader.GetModel().CreateEntry()
+        reader.LoadEntry(0, entry)
+        self.assertEqual(entry["f"], 42)