From 95a8fa5dea6c9b4d765789ac4baa9fbd5902ab44 Mon Sep 17 00:00:00 2001
From: Mateusz Jakub Fila <37295697+m-fila@users.noreply.github.com>
Date: Thu, 12 Sep 2024 13:30:43 +0200
Subject: [PATCH] make `CreateDataFrame` accessible from python (#674)

* fix CreateDataFrame visibility in python

* use load dict to read datasource lib in python

* fix pylint warnings
---
 python/podio/__init__.py         |  8 +++++++-
 python/podio/data_source.py      | 10 ++++++++++
 src/rds_selection.xml            |  1 +
 tests/root_io/CMakeLists.txt     |  6 ++++++
 tests/root_io/read_datasource.py | 12 ++++++++++++
 5 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 python/podio/data_source.py
 create mode 100644 tests/root_io/read_datasource.py

diff --git a/python/podio/__init__.py b/python/podio/__init__.py
index 5baf354f1..147dc96c0 100644
--- a/python/podio/__init__.py
+++ b/python/podio/__init__.py
@@ -21,4 +21,10 @@
 except ImportError:
     pass
 
-__all__ = ["__version__", "Frame", "root_io", "sio_io", "reading", "version"]
+try:
+    # Same mechanism as for the sio_io above
+    from . import data_source
+except ImportError:
+    pass
+
+__all__ = ["__version__", "Frame", "root_io", "sio_io", "reading", "data_source", "version"]
diff --git a/python/podio/data_source.py b/python/podio/data_source.py
new file mode 100644
index 000000000..cead7b890
--- /dev/null
+++ b/python/podio/data_source.py
@@ -0,0 +1,10 @@
+"""Python module for creating ROOT RDataFrame with files containing podio Frames"""
+
+from ROOT import gSystem
+
+if gSystem.Load("libpodioDataSourceDict") < 0:
+    raise ImportError("Error when loading libpodioDataSourceDict")
+
+from ROOT import podio  # pylint: disable=wrong-import-position
+
+CreateDataFrame = podio.CreateDataFrame
diff --git a/src/rds_selection.xml b/src/rds_selection.xml
index ead05ceac..a4b2b7f4b 100644
--- a/src/rds_selection.xml
+++ b/src/rds_selection.xml
@@ -1,5 +1,6 @@
 <lcgdict>
   <selection>
     <class name="podio::DataSource"/>
+    <function name="podio::CreateDataFrame"/>
   </selection>
 </lcgdict>
diff --git a/tests/root_io/CMakeLists.txt b/tests/root_io/CMakeLists.txt
index 9d97ab916..9de085a4f 100644
--- a/tests/root_io/CMakeLists.txt
+++ b/tests/root_io/CMakeLists.txt
@@ -79,3 +79,9 @@ endif()
 add_test(NAME param_reading_rdataframe COMMAND python3 ${CMAKE_CURRENT_SOURCE_DIR}/param_reading_rdataframe.py example_frame.root)
 PODIO_SET_TEST_ENV(param_reading_rdataframe)
 set_property(TEST param_reading_rdataframe PROPERTY DEPENDS write_frame_root)
+
+if(ENABLE_DATASOURCE)
+  add_test(NAME read_python_with_rdatasource_root COMMAND python3 ${PROJECT_SOURCE_DIR}/tests/root_io/read_datasource.py)
+  PODIO_SET_TEST_ENV(read_python_with_rdatasource_root)
+  set_property(TEST read_python_with_rdatasource_root PROPERTY DEPENDS read_with_rdatasource_root)
+endif()
diff --git a/tests/root_io/read_datasource.py b/tests/root_io/read_datasource.py
new file mode 100644
index 000000000..4b6318dac
--- /dev/null
+++ b/tests/root_io/read_datasource.py
@@ -0,0 +1,12 @@
+#!/usr/bin/env python3
+"""Small test case for checking DataSource based creating RDataFrames is accessible from python"""
+
+import ROOT
+from podio.data_source import CreateDataFrame  # pylint: disable=import-error, no-name-in-module
+
+if ROOT.gSystem.Load("libTestDataModelDict") < 0:
+    raise RuntimeError("Could not load TestDataModel dictionary")
+
+rdf = CreateDataFrame("example_frame.root")
+
+assert rdf.Count().GetValue() == 10