[6.42] [rfile][python] add support for __getitem__ and __contains__

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py

@@ -24,6 +24,23 @@
     # file.Flush()
 \endcode
 
+Getting an object can also be done via the dictionary syntax:
+
+\code{.py}
+myHist = file['myHisto']
+\endcode
+
+The difference between the two methods is that Get will return `None` if the object is missing, while the square
+bracket operator will raise a `KeyError`.
+
+There are two main ways to test for the presence of objects in the file:
+
+\code{.py}
+histoIsThere = 'myHisto' in file
+# or
+histoIsThere = file.GetKeyInfo('myHisto') is None
+\endcode
+
 The Put method will raise an error if an object is already present at the given path; if you want to overwrite an
 existing object use the Overwrite method instead:
 
@@ -172,7 +189,7 @@ class _RFile_Overwrite:
     def __init__(self, rfile):
         self._rfile = rfile
 
-    def __call__(self, path, obj = None, title = "", *, createNewCycle = True):
+    def __call__(self, path, obj=None, title="", *, createNewCycle=True):
         """
         Non-templated Overwrite()
         """
@@ -187,6 +204,7 @@ def __call__(self, path, obj = None, title = "", *, createNewCycle = True):
             className = objType.__cpp_name__
 
         from ROOT import RFile
+
         flags = (0 if createNewCycle else 1) * RFile.kOverwrite_NoNewCycle
         self._rfile.Overwrite[className](path, obj, title, flags)
 
@@ -232,14 +250,14 @@ def _RFileInit(rfile):
     raise NotImplementedError("RFile can only be created via Recreate, Open or Update")
 
 
-def _GetKeyInfo(rfile, path):
+def _RFile_GetKeyInfo(rfile, path):
     key = rfile._OriginalGetKeyInfo(path)
     if key.has_value():
         return key.value()
     return None
 
 
-def _ListKeys(rfile, basePath="", **kwargs):
+def _RFile_ListKeys(rfile, basePath="", **kwargs):
     """
     Returns an iterable over all keys of objects and/or directories written into this RFile starting at path
     `basePath` (defaulting to include the content of all subdirectories).
@@ -273,15 +291,27 @@ def _ListKeys(rfile, basePath="", **kwargs):
     """
     from ROOT import RFile
 
-    listObjects = kwargs['listObjects'] if 'listObjects' in kwargs else True
-    listDirs = kwargs['listDirs'] if 'listDirs' in kwargs else False
-    listRecursive = kwargs['listRecursive'] if 'listRecursive' in kwargs else True
+    listObjects = kwargs["listObjects"] if "listObjects" in kwargs else True
+    listDirs = kwargs["listDirs"] if "listDirs" in kwargs else False
+    listRecursive = kwargs["listRecursive"] if "listRecursive" in kwargs else True
 
     flags = (listObjects * RFile.kListObjects) | (listDirs * RFile.kListDirs) | (listRecursive * RFile.kListRecursive)
     iter = rfile._OriginalListKeys(basePath, flags)
     return iter
 
 
+def _RFile_getitem(rfile, namecycle):
+    item = rfile.Get(namecycle)
+    if item is None:
+        raise KeyError(namecycle)
+    return item
+
+
+def _RFile_contains(rfile, namecycle):
+    keyInfo = _RFile_GetKeyInfo(rfile, namecycle)
+    return keyInfo is not None
+
+
 @pythonization("RFile", ns="ROOT")
 def pythonize_rfile(klass):
     # Explicitly prevent to create a RFile via ctor
@@ -297,6 +327,8 @@ def pythonize_rfile(klass):
     klass.__enter__ = lambda rfile: rfile
     klass.__exit__ = _RFileExit
     klass._OriginalGetKeyInfo = klass.GetKeyInfo
-    klass.GetKeyInfo = _GetKeyInfo
+    klass.GetKeyInfo = _RFile_GetKeyInfo
     klass._OriginalListKeys = klass.ListKeys
-    klass.ListKeys = _ListKeys
+    klass.ListKeys = _RFile_ListKeys
+    klass.__getitem__ = _RFile_getitem
+    klass.__contains__ = _RFile_contains

io/io/test/rfile.py

@@ -8,8 +8,11 @@
 def platform_supports_assertRaises():
     # TODO: re-enable it on Windows once the exception handling is fixed
     # TODO: remove mac 15 exception once we stop supporting it
-    return platform.system() != "Windows" and (platform.system() != "Darwin" or not platform.mac_ver()[0].startswith('15'))
-
+    return platform.system() != "Windows" and (
+        platform.system() != "Darwin" or not platform.mac_ver()[0].startswith("15")
+    )
+
+
 class RFileTests(unittest.TestCase):
     def test_open_for_reading(self):
         """A RFile can read a ROOT file created by TFile"""
@@ -81,7 +84,6 @@ def test_getkeyinfo(self):
         finally:
             os.remove(fileName)
 
-
     def test_listkeys(self):
         """A RFile can query the keys of its objects and directories"""
 
@@ -108,17 +110,25 @@ def test_listkeys(self):
                     ["foo/hist", "foo/bar/hist", "foo/bar/hist2", "foo/hist2"],
                 )
 
-                self.assertEqual([key.GetPath() for key in rfile.ListKeys("foo/bar")], ["foo/bar/hist", "foo/bar/hist2"])
+                self.assertEqual(
+                    [key.GetPath() for key in rfile.ListKeys("foo/bar")], ["foo/bar/hist", "foo/bar/hist2"]
+                )
 
                 self.assertEqual(
                     [key.GetPath() for key in rfile.ListKeys("", listDirs=True, listObjects=False)], ["foo", "foo/bar"]
                 )
                 self.assertEqual(
-                    [key.GetPath() for key in rfile.ListKeys("", listDirs=True, listObjects=False, listRecursive=False)],
+                    [
+                        key.GetPath()
+                        for key in rfile.ListKeys("", listDirs=True, listObjects=False, listRecursive=False)
+                    ],
                     ["foo"],
                 )
                 self.assertEqual(
-                    [key.GetPath() for key in rfile.ListKeys("foo", listDirs=True, listObjects=False, listRecursive=False)],
+                    [
+                        key.GetPath()
+                        for key in rfile.ListKeys("foo", listDirs=True, listObjects=False, listRecursive=False)
+                    ],
                     ["foo/bar"],
                 )
                 self.assertEqual(
@@ -236,8 +246,46 @@ def test_overwrite(self):
                 self.assertTrue(isinstance(hist, ROOT.TH1F))
 
         finally:
-            # os.remove(fileName)
-            pass
+            os.remove(fileName)
+
+    def test_getitem(self):
+        """Objects may be retrieved from RFile via the [] operator. KeyError is thrown if the object isn't there."""
+
+        fileName = "test_rfile_getitem_py.root"
+        try:
+            with ROOT.RFile.Recreate(fileName) as rfile:
+                hist = ROOT.TH1D("hist", "", 100, -10, 10)
+                hist.FillRandom("gaus", 10)
+                rfile.Put("hist", hist)
+                # There is no setitem support.
+                if platform_supports_assertRaises():
+                    with self.assertRaisesRegex(TypeError, "does not support item assignment"):
+                        rfile["hist2"] = hist
+
+            with ROOT.RFile.Open(fileName) as rfile:
+                hist = rfile["hist"]
+                self.assertNotEqual(hist, None)
+                if platform_supports_assertRaises():
+                    with self.assertRaises(KeyError):
+                        _hist2 = rfile["hist2"]
+
+        finally:
+            os.remove(fileName)
+
+    def test_contains(self):
+        """Can query presence of objects via `'key' in rfile`"""
+
+        fileName = "test_rfile_contains_py.root"
+        try:
+            with ROOT.RFile.Recreate(fileName) as rfile:
+                hist = ROOT.TH1D("hist", "", 100, -10, 10)
+                hist.FillRandom("gaus", 10)
+                rfile.Put("hist", hist)
+                self.assertTrue("hist" in rfile)
+                self.assertFalse("hist2" in rfile)
+
+        finally:
+            os.remove(fileName)
 
 
 if __name__ == "__main__":