diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py index d1e27e640d2ab..727835e1618ea 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_rfile.py @@ -66,8 +66,8 @@ def __call__(self, name, obj): if isinstance(obj, str): # special case: automatically convert python str to std::string className = "std::string" - elif not hasattr(objType, '__cpp_name__'): - raise TypeError(f"type {objType} is not supported by ROOT I/O") + elif not hasattr(objType, "__cpp_name__"): + raise TypeError(f"type {objType} is not supported by ROOT I/O") else: className = objType.__cpp_name__ self._rfile.Put[className](name, obj) @@ -119,9 +119,44 @@ def _GetKeyInfo(rfile, path): return None -def _ListKeys(rfile, basePath="", listObjects = True, listDirs = False, listRecursive = True): +def _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). + By default, keys referring to directories are not returned: only those referring to leaf objects are. + If `basePath` is the path of a leaf object, only `basePath` itself will be returned. + If it is the path of a directory, it will not be included in the listing. + + You can specify what to list via the keyword arguments: + - if `listObjects == True`, the listing will include keys of non-directory objects (default); + - if `listDirs == True`, the listing will include keys of directory objects; + - if `listRecursive == True`, the listing will recurse on all subdirectories of `basePath` (default), + otherwise it will only list immediate children of `basePath`. + Example usage: + ~~~{.py} + for key in file.ListKeys(): + # iterate over all objects in the RFile + print(f"{key.GetPath()};{key.GetCycle()} of type {key.GetClassName()}") + + for key in file.ListKeys("", listDirs=True): + # iterate over all objects and directories in the RFile + print(f"{key.GetPath()};{key.GetCycle()} of type {key.GetClassName()}") + + for key in file.ListKeys("a/b", listRecursive=False): + # iterate over all objects that are immediate children of directory "a/b" + print(f"{key.GetPath()};{key.GetCycle()} of type {key.GetClassName()}") + + for key in file.ListKeys("foo", listDirs=True, listObjects=False): + # iterate over all directories under directory "foo", recursively + print(key.GetPath()) + ~~~ + """ from ROOT.Experimental 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 + flags = (listObjects * RFile.kListObjects) | (listDirs * RFile.kListDirs) | (listRecursive * RFile.kListRecursive) iter = rfile._OriginalListKeys(basePath, flags) return iter diff --git a/io/io/inc/ROOT/RFile.hxx b/io/io/inc/ROOT/RFile.hxx index c08250a2d752b..40bd0a2ff31fd 100644 --- a/io/io/inc/ROOT/RFile.hxx +++ b/io/io/inc/ROOT/RFile.hxx @@ -376,6 +376,7 @@ public: /// `basePath` (defaulting to include the content of all subdirectories). /// By default, keys referring to directories are not returned: only those referring to leaf objects are. /// If `basePath` is the path of a leaf object, only `basePath` itself will be returned. + /// If `basePath` is the path of a directory, it won't appear in the listing. /// `flags` is a bitmask specifying the listing mode. /// If `(flags & kListObjects) != 0`, the listing will include keys of non-directory objects (default); /// If `(flags & kListDirs) != 0`, the listing will include keys of directory objects; diff --git a/io/io/src/RFile.cxx b/io/io/src/RFile.cxx index ca3a979fe6795..bc80ffa04170c 100644 --- a/io/io/src/RFile.cxx +++ b/io/io/src/RFile.cxx @@ -475,6 +475,10 @@ void ROOT::Experimental::RFileKeyIterable::RIterator::Advance() assert(!fIterStack.empty()); const std::size_t nesting = fIterStack.size() - 1; + // Skip key if it's the root dir itself + if (isDir && fullPath == fPattern) + continue; + // Skip key if it's not a child of root dir if (!ROOT::StartsWith(fullPath, fPattern)) continue; diff --git a/io/io/test/rfile.cxx b/io/io/test/rfile.cxx index e08d7b2c14751..14e45e5be2fd9 100644 --- a/io/io/test/rfile.cxx +++ b/io/io/test/rfile.cxx @@ -482,10 +482,10 @@ TEST(RFile, IterateKeysOnlyDirs) { auto file = RFile::Open(fileGuard.GetPath()); EXPECT_EQ(JoinKeyNames(file->ListKeys("", RFile::kListDirs | RFile::kListRecursive)), "a, a/b, e, e/c"); - EXPECT_EQ(JoinKeyNames(file->ListKeys("a", RFile::kListDirs | RFile::kListRecursive)), "a, a/b"); - EXPECT_EQ(JoinKeyNames(file->ListKeys("a/b", RFile::kListDirs | RFile::kListRecursive)), "a/b"); + EXPECT_EQ(JoinKeyNames(file->ListKeys("a", RFile::kListDirs | RFile::kListRecursive)), "a/b"); + EXPECT_EQ(JoinKeyNames(file->ListKeys("a/b", RFile::kListDirs | RFile::kListRecursive)), ""); EXPECT_EQ(JoinKeyNames(file->ListKeys("a/b/c", RFile::kListDirs | RFile::kListRecursive)), ""); - EXPECT_EQ(JoinKeyNames(file->ListKeys("e", RFile::kListDirs | RFile::kListRecursive)), "e, e/c"); + EXPECT_EQ(JoinKeyNames(file->ListKeys("e", RFile::kListDirs | RFile::kListRecursive)), "e/c"); } } @@ -506,10 +506,10 @@ TEST(RFile, IterateKeysOnlyDirsNonRecursive) { auto file = RFile::Open(fileGuard.GetPath()); EXPECT_EQ(JoinKeyNames(file->ListKeys("", RFile::kListDirs)), "a, e"); - EXPECT_EQ(JoinKeyNames(file->ListKeys("a", RFile::kListDirs)), "a, a/b"); - EXPECT_EQ(JoinKeyNames(file->ListKeys("a/b", RFile::kListDirs)), "a/b"); + EXPECT_EQ(JoinKeyNames(file->ListKeys("a", RFile::kListDirs)), "a/b"); + EXPECT_EQ(JoinKeyNames(file->ListKeys("a/b", RFile::kListDirs)), ""); EXPECT_EQ(JoinKeyNames(file->ListKeys("a/b/c", RFile::kListDirs)), ""); - EXPECT_EQ(JoinKeyNames(file->ListKeys("e", RFile::kListDirs)), "e, e/c"); + EXPECT_EQ(JoinKeyNames(file->ListKeys("e", RFile::kListDirs)), "e/c"); } } diff --git a/io/io/test/rfile.py b/io/io/test/rfile.py index 4dfd06e8d0d4e..30d78872665b1 100644 --- a/io/io/test/rfile.py +++ b/io/io/test/rfile.py @@ -115,6 +115,10 @@ def test_listkeys(self): [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)], + ["foo/bar"], + ) self.assertEqual( [key.GetPath() for key in rfile.ListKeys("", listDirs=True, listRecursive=False)], ["hist", "foo"] )