Skip to content

Commit

Permalink
Add compatibility with pathlib.Path to python readers (#699)
Browse files Browse the repository at this point in the history
* add compatibility with pathlib.Path to root_io

* outsource TypeError handling to os.fspath

* add pathlib compatibility in sio_io as well

* sio_io: only store first path in list returned by convert_to_str_paths

* make ruff import sorting happy

* outsource convert_to_str_paths

* Fix imports to be absolute

* Add module docstring to fix pylint complaints
  • Loading branch information
Victor-Schwan authored Nov 8, 2024
1 parent b5cddb7 commit d1bb7b0
Show file tree
Hide file tree
Showing 3 changed files with 45 additions and 19 deletions.
25 changes: 11 additions & 14 deletions python/podio/root_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@

from podio.base_reader import BaseReaderMixin # pylint: disable=wrong-import-position # noqa: E402
from podio.base_writer import BaseWriterMixin # pylint: disable=wrong-import-position # noqa: E402
from podio.utils import convert_to_str_paths # pylint: disable=wrong-import-position # noqa: E402


class Reader(BaseReaderMixin):
Expand All @@ -17,11 +18,9 @@ def __init__(self, filenames):
"""Create a reader that reads from the passed file(s).
Args:
filenames (str or list[str]): file(s) to open and read data from
filenames (str or list[str] or Path or list[Path]): file(s) to open and read data from
"""
if isinstance(filenames, str):
filenames = (filenames,)

filenames = convert_to_str_paths(filenames)
self._reader = podio.ROOTReader()
self._reader.openFiles(filenames)

Expand All @@ -35,11 +34,9 @@ def __init__(self, filenames):
"""Create an RNTuple reader that reads from the passed file(s).
Args:
filenames (str or list[str]): file(s) to open and read data from
filenames (str or list[str] or Path or list[Path]): file(s) to open and read data from
"""
if isinstance(filenames, str):
filenames = (filenames,)

filenames = convert_to_str_paths(filenames)
self._reader = podio.RNTupleReader()
self._reader.openFiles(filenames)

Expand All @@ -57,11 +54,9 @@ def __init__(self, filenames):
"""Create a reader that reads from the passed file(s).
Args:
filenames (str or list[str]): file(s) to open and read data from
filenames (str or list[str] or Path or list[Path]): file(s) to open and read data from
"""
if isinstance(filenames, str):
filenames = (filenames,)

filenames = convert_to_str_paths(filenames)
self._reader = podio.ROOTLegacyReader()
self._reader.openFiles(filenames)
self._is_legacy = True
Expand All @@ -76,8 +71,9 @@ def __init__(self, filename):
"""Create a writer for writing files
Args:
filename (str): The name of the output file
filename (str or Path): The name of the output file
"""
filename = convert_to_str_paths(filename)[0]
self._writer = podio.ROOTWriter(filename)
super().__init__()

Expand All @@ -89,7 +85,8 @@ def __init__(self, filename):
"""Create a writer for writing files
Args:
filename (str): The name of the output file
filename (str or Path): The name of the output file
"""
filename = convert_to_str_paths(filename)[0]
self._writer = podio.RNTupleWriter(filename)
super().__init__()
15 changes: 10 additions & 5 deletions python/podio/sio_io.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@

from podio.base_reader import BaseReaderMixin # pylint: disable=wrong-import-position
from podio.base_writer import BaseWriterMixin # pylint: disable=wrong-import-position
from podio.utils import convert_to_str_paths # pylint: disable=wrong-import-position # noqa: E402


class Reader(BaseReaderMixin):
Expand All @@ -20,8 +21,9 @@ def __init__(self, filename):
"""Create a reader that reads from the passed file.
Args:
filename (str): File to open and read data from
filename (str or Path): File to open and read data from.
"""
filename = convert_to_str_paths(filename)[0]
self._reader = podio.SIOReader()
self._reader.openFile(filename)

Expand All @@ -39,8 +41,9 @@ def __init__(self, filename):
"""Create a reader that reads from the passed file.
Args:
filename (str): File to open and read data from
filename (str or Path): File to open and read data from.
"""
filename = convert_to_str_paths(filename)[0]
self._reader = podio.SIOLegacyReader()
self._reader.openFile(filename)
self._is_legacy = True
Expand All @@ -49,13 +52,15 @@ def __init__(self, filename):


class Writer(BaseWriterMixin):
"""Writer class for writing podio root files"""
"""Writer class for writing podio root files."""

def __init__(self, filename):
"""Create a writer for writing files
"""Create a writer for writing files.
Args:
filename (str): The name of the output file
filename (str or Path): The name of the output file.
"""
filename = convert_to_str_paths(filename)[0]
self._writer = podio.SIOWriter(filename)

super().__init__()
24 changes: 24 additions & 0 deletions python/podio/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
#!/usr/bin/env python3
"""utility functionality for podio"""

import os
from collections.abc import Iterable
from pathlib import Path


def convert_to_str_paths(filenames):
"""Converts filenames to string paths, handling both string and pathlib.Path objects and
iterables of such objects.
Args:
filenames (str, Path, or Iterable[str | Path]): A single filepath or an iterable of
filepaths to convert to str object(s).
Returns:
list[str]: A list of filepaths as strings.
"""

if isinstance(filenames, Iterable) and not isinstance(filenames, (str, Path)):
return [os.fspath(fn) for fn in filenames]

return [os.fspath(filenames)]

0 comments on commit d1bb7b0

Please sign in to comment.