From e56d9567f776e52d09ef5ff418ed4b49c0832c6e Mon Sep 17 00:00:00 2001
From: Mateusz Jakub Fila <mateusz.jakub.fila@cern.ch>
Date: Wed, 11 Sep 2024 17:25:54 +0200
Subject: [PATCH] fix CreateDataFrame visibility in python

---
 python/podio/__init__.py         |  8 +++++++-
 python/podio/data_source.py      | 13 +++++++++++++
 tests/root_io/CMakeLists.txt     |  6 ++++++
 tests/root_io/read_datasource.py | 13 +++++++++++++
 4 files changed, 39 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..5556a7388
--- /dev/null
+++ b/python/podio/data_source.py
@@ -0,0 +1,13 @@
+"""Python module for creating ROOT RDataFrame with files containing podio Frames"""
+
+from ROOT import gSystem, gInterpreter
+
+if not gSystem.DynamicPathName("libpodioDataSource", True):
+    raise ImportError("Error finding libpodioDataSource")
+
+if gInterpreter.LoadFile("podio/DataSource.h") != 0:
+    raise ImportError("Error when loading file podio/DataSource.h")
+
+from ROOT import podio  # pylint: disable=wrong-import-position
+
+CreateDataFrame = podio.CreateDataFrame
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..40b0b5c93
--- /dev/null
+++ b/tests/root_io/read_datasource.py
@@ -0,0 +1,13 @@
+#!/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
+
+if ROOT.gSystem.Load("libTestDataModelDict") < 0:  # noqa: E402
+    raise RuntimeError("Could not load TestDataModel dictionary")
+
+input_file = "example_frame.root"  # pylint: disable-msg=C0103
+rdf = CreateDataFrame(input_file)
+
+assert rdf.Count().GetValue() == 10