Skip to content

Commit

Permalink
expose dataloaders to python SDK
Browse files Browse the repository at this point in the history
  • Loading branch information
teh-cmc committed Feb 29, 2024
1 parent ffe0a30 commit 5dd9685
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 1 deletion.
6 changes: 5 additions & 1 deletion rerun_py/rerun_sdk/rerun/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,8 +77,10 @@
"get_recording_id",
"get_thread_local_data_recording",
"is_enabled",
"log_components",
"log",
"log_components",
"log_file_from_contents",
"log_file_from_path",
"memory_recording",
"new_entity_path",
"reset_time",
Expand All @@ -105,6 +107,8 @@
escape_entity_path_part,
log,
log_components,
log_file_from_contents,
log_file_from_path,
new_entity_path,
)
from .any_value import AnyValues
Expand Down
67 changes: 67 additions & 0 deletions rerun_py/rerun_sdk/rerun/_log.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

from pathlib import Path
from typing import Any, Iterable

import pyarrow as pa
Expand Down Expand Up @@ -282,6 +283,72 @@ def log_components(
)


@catch_and_log_exceptions()
def log_file_from_path(
file_path: str | Path,
*,
recording: RecordingStream | None = None,
) -> None:
r"""
Logs the file at the given `path` using all `re_data_source::DataLoader`s available.
A single `path` might be handled by more than one loader.
This method blocks until either at least one `re_data_source::DataLoader` starts
streaming data in or all of them fail.
See <https://www.rerun.io/docs/howto/open-any-file> for more information.
Parameters
----------
file_path:
Path to the file to be logged.
recording:
Specifies the [`rerun.RecordingStream`][] to use. If left unspecified,
defaults to the current active data recording, if there is one. See
also: [`rerun.init`][], [`rerun.set_global_data_recording`][].
"""

bindings.log_file_from_path(Path(file_path), recording=recording)


@catch_and_log_exceptions()
def log_file_from_contents(
file_path: str | Path,
file_contents: bytes,
*,
recording: RecordingStream | None = None,
) -> None:
r"""
Logs the given `file_contents` using all `crate::DataLoader`s available.
A single `path` might be handled by more than one loader.
This method blocks until either at least one `re_data_source::DataLoader` starts
streaming data in or all of them fail.
See <https://www.rerun.io/docs/howto/open-any-file> for more information.
Parameters
----------
file_path:
Path to the file that the `file_contents` belong to.
file_contents:
Contents to be logged.
recording:
Specifies the [`rerun.RecordingStream`][] to use. If left unspecified,
defaults to the current active data recording, if there is one. See
also: [`rerun.init`][], [`rerun.set_global_data_recording`][].
"""

bindings.log_file_from_contents(Path(file_path), file_contents, recording=recording)


def escape_entity_path_part(part: str) -> str:
r"""
Escape an individual part of an entity path.
Expand Down
50 changes: 50 additions & 0 deletions rerun_py/src/python_bridge.rs
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,8 @@ fn rerun_bindings(_py: Python<'_>, m: &PyModule) -> PyResult<()> {

// log any
m.add_function(wrap_pyfunction!(log_arrow_msg, m)?)?;
m.add_function(wrap_pyfunction!(log_file_from_path, m)?)?;
m.add_function(wrap_pyfunction!(log_file_from_contents, m)?)?;

// misc
m.add_function(wrap_pyfunction!(version, m)?)?;
Expand Down Expand Up @@ -962,6 +964,54 @@ fn log_arrow_msg(
Ok(())
}

#[pyfunction]
#[pyo3(signature = (
file_path,
recording=None,
))]
fn log_file_from_path(
py: Python<'_>,
file_path: std::path::PathBuf,
recording: Option<&PyRecordingStream>,
) -> PyResult<()> {
let Some(recording) = get_data_recording(recording) else {
return Ok(());
};

recording
.log_file_from_path(file_path)
.map_err(|err| PyRuntimeError::new_err(err.to_string()))?;

py.allow_threads(flush_garbage_queue);

Ok(())
}

#[pyfunction]
#[pyo3(signature = (
file_path,
file_contents,
recording=None,
))]
fn log_file_from_contents(
py: Python<'_>,
file_path: std::path::PathBuf,
file_contents: &[u8],
recording: Option<&PyRecordingStream>,
) -> PyResult<()> {
let Some(recording) = get_data_recording(recording) else {
return Ok(());
};

recording
.log_file_from_contents(file_path, std::borrow::Cow::Borrowed(file_contents))
.map_err(|err| PyRuntimeError::new_err(err.to_string()))?;

py.allow_threads(flush_garbage_queue);

Ok(())
}

// --- Misc ---

/// Return a verbose version string
Expand Down

0 comments on commit 5dd9685

Please sign in to comment.