Skip to content

Commit

Permalink
make CreateDataFrame accessible from python (#674)
Browse files Browse the repository at this point in the history
* fix CreateDataFrame visibility in python

* use load dict to read datasource lib in python

* fix pylint warnings
  • Loading branch information
m-fila authored Sep 12, 2024
1 parent cf1f33c commit 95a8fa5
Show file tree
Hide file tree
Showing 5 changed files with 36 additions and 1 deletion.
8 changes: 7 additions & 1 deletion python/podio/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
10 changes: 10 additions & 0 deletions python/podio/data_source.py
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions src/rds_selection.xml
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
<lcgdict>
<selection>
<class name="podio::DataSource"/>
<function name="podio::CreateDataFrame"/>
</selection>
</lcgdict>
6 changes: 6 additions & 0 deletions tests/root_io/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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()
12 changes: 12 additions & 0 deletions tests/root_io/read_datasource.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit 95a8fa5

Please sign in to comment.