expose dataloaders to python SDK

teh-cmc · teh-cmc · commit 5dd968596144 · 2024-02-29T13:03:41.000+01:00
diff --git a/rerun_py/rerun_sdk/rerun/__init__.py b/rerun_py/rerun_sdk/rerun/__init__.py
@@ -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",
@@ -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
diff --git a/rerun_py/rerun_sdk/rerun/_log.py b/rerun_py/rerun_sdk/rerun/_log.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Any, Iterable
 
 import pyarrow as pa
@@ -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.
diff --git a/rerun_py/src/python_bridge.rs b/rerun_py/src/python_bridge.rs
@@ -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)?)?;
@@ -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
