DataLoaders 7: support for custom DataLoaders

crates/re_data_source/src/data_loader/mod.rs

@@ -229,8 +229,26 @@ static BUILTIN_LOADERS: Lazy<Vec<Arc<dyn DataLoader>>> = Lazy::new(|| {
 
 /// Iterator over all registered [`DataLoader`]s.
 #[inline]
-pub fn iter_loaders() -> impl ExactSizeIterator<Item = Arc<dyn DataLoader>> {
-    BUILTIN_LOADERS.clone().into_iter()
+pub fn iter_loaders() -> impl Iterator<Item = Arc<dyn DataLoader>> {
+    BUILTIN_LOADERS
+        .clone()
+        .into_iter()
+        .chain(CUSTOM_LOADERS.read().clone())
+}
+
+/// Keeps track of all custom [`DataLoader`]s.
+///
+/// Use [`register_custom_data_loader`] to add new loaders.
+static CUSTOM_LOADERS: Lazy<parking_lot::RwLock<Vec<Arc<dyn DataLoader>>>> =
+    Lazy::new(parking_lot::RwLock::default);
+
+/// Register a custom [`DataLoader`].
+///
+/// Any time the Rerun Viewer opens a file or directory, this custom loader will be notified.
+/// Refer to [`DataLoader`]'s documentation for more information.
+#[inline]
+pub fn register_custom_data_loader(loader: impl DataLoader + 'static) {
+    CUSTOM_LOADERS.write().push(Arc::new(loader));
 }
 
 // ---

crates/re_data_source/src/lib.rs

@@ -15,8 +15,8 @@ mod web_sockets;
 mod load_stdin;
 
 pub use self::data_loader::{
-    iter_loaders, ArchetypeLoader, DataLoader, DataLoaderError, DirectoryLoader, LoadedData,
-    RrdLoader,
+    iter_loaders, register_custom_data_loader, ArchetypeLoader, DataLoader, DataLoaderError,
+    DirectoryLoader, LoadedData, RrdLoader,
 };
 pub use self::data_source::DataSource;
 pub use self::load_file::{extension, load_from_file_contents};

crates/re_data_source/src/load_file.rs

@@ -34,14 +34,16 @@ pub fn load_from_path(
 
     re_log::info!("Loading {path:?}…");
 
-    let store_info = prepare_store_info(store_id, file_source, path, path.is_dir());
+    let data = load(store_id, path, None)?;
+
+    // If we reach this point, then at least one compatible `DataLoader` has been found.
+    let store_info = prepare_store_info(store_id, file_source, path);
     if let Some(store_info) = store_info {
         if tx.send(store_info).is_err() {
             return Ok(()); // other end has hung up.
         }
     }
 
-    let data = load(store_id, path, None)?;
     send(store_id, data, tx);
 
     Ok(())
@@ -67,14 +69,16 @@ pub fn load_from_file_contents(
 
     re_log::info!("Loading {filepath:?}…");
 
-    let store_info = prepare_store_info(store_id, file_source, filepath, false);
+    let data = load(store_id, filepath, Some(contents))?;
+
+    // If we reach this point, then at least one compatible `DataLoader` has been found.
+    let store_info = prepare_store_info(store_id, file_source, filepath);
     if let Some(store_info) = store_info {
         if tx.send(store_info).is_err() {
             return Ok(()); // other end has hung up.
         }
     }
 
-    let data = load(store_id, filepath, Some(contents))?;
     send(store_id, data, tx);
 
     Ok(())
@@ -92,20 +96,11 @@ pub fn extension(path: &std::path::Path) -> String {
         .to_string()
 }
 
-/// Returns whether the given path is supported by builtin [`crate::DataLoader`]s.
-///
-/// This does _not_ access the filesystem.
-#[inline]
-pub fn is_associated_with_builtin_loader(path: &std::path::Path, is_dir: bool) -> bool {
-    is_dir || crate::is_supported_file_extension(&extension(path))
-}
-
 /// Prepares an adequate [`re_log_types::StoreInfo`] [`LogMsg`] given the input.
 pub(crate) fn prepare_store_info(
     store_id: &re_log_types::StoreId,
     file_source: FileSource,
     path: &std::path::Path,
-    is_dir: bool,
 ) -> Option<LogMsg> {
     re_tracing::profile_function!(path.display().to_string());
 
@@ -114,10 +109,9 @@ pub(crate) fn prepare_store_info(
     let app_id = re_log_types::ApplicationId(path.display().to_string());
     let store_source = re_log_types::StoreSource::File { file_source };
 
-    let is_builtin = is_associated_with_builtin_loader(path, is_dir);
     let is_rrd = crate::SUPPORTED_RERUN_EXTENSIONS.contains(&extension(path).as_str());
 
-    (!is_rrd && is_builtin).then(|| {
+    (!is_rrd).then(|| {
         LogMsg::SetStoreInfo(SetStoreInfo {
             row_id: re_log_types::RowId::new(),
             info: re_log_types::StoreInfo {

examples/rust/custom_data_loader/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "custom_data_loader"
+version = "0.12.0-alpha.1+dev"
+edition = "2021"
+rust-version = "1.72"
+license = "MIT OR Apache-2.0"
+publish = false
+
+[dependencies]
+rerun = { path = "../../../crates/rerun", features = ["native_viewer"] }
+
+[build-dependencies]
+re_build_tools = { path = "../../../crates/re_build_tools" }

examples/rust/custom_data_loader/README.md

@@ -0,0 +1,20 @@
+---
+title: Custom data-loader
+tags: [data-loader, extension]
+rust: https://github.com/rerun-io/rerun/tree/latest/examples/rust/custom_data_loader/src/main.rs?speculative-link
+---
+
+<picture>
+  <img src="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/full.png" alt="">
+  <source media="(max-width: 480px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/480w.png">
+  <source media="(max-width: 768px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/768w.png">
+  <source media="(max-width: 1024px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/1024w.png">
+  <source media="(max-width: 1200px)" srcset="https://static.rerun.io/custom_data_loader/e44aadfa02fade5a3cf5d7cbdd6e0bf65d9f6446/1200w.png">
+</picture>
+
+This example demonstrates how to implement and register a `DataLoader` into the Rerun Viewer in order to add support for loading arbitrary files.
+
+Usage:
+```sh
+$ cargo r -p custom_data_loader -- path/to/some/file
+```

examples/rust/custom_data_loader/build.rs

@@ -0,0 +1,3 @@
+fn main() {
+    re_build_tools::export_build_info_vars_for_crate("custom_data_loader");
+}

examples/rust/custom_data_loader/src/main.rs

@@ -0,0 +1,82 @@
+//! This example demonstrates how to implement and register a [`re_data_source::DataLoader`] into
+//! the Rerun Viewer in order to add support for loading arbitrary files.
+//!
+//! Usage:
+//! ```sh
+//! $ cargo r -p custom_data_loader -- path/to/some/file
+//! ```
+
+use rerun::{
+    external::{anyhow, re_build_info, re_data_source, re_log, tokio},
+    log::{DataRow, RowId},
+    EntityPath, TimePoint,
+};
+
+#[tokio::main]
+async fn main() -> anyhow::Result<std::process::ExitCode> {
+    re_log::setup_native_logging();
+
+    re_data_source::register_custom_data_loader(HashLoader);
+
+    let build_info = re_build_info::build_info!();
+    rerun::run(build_info, rerun::CallSource::Cli, std::env::args())
+        .await
+        .map(std::process::ExitCode::from)
+}
+
+// ---
+
+/// A custom [`re_data_source::DataLoader`] that logs the hash of file as a [`rerun::TextDocument`].
+struct HashLoader;
+
+impl re_data_source::DataLoader for HashLoader {
+    fn name(&self) -> String {
+        "rerun.data_loaders.HashLoader".into()
+    }
+
+    fn load_from_path(
+        &self,
+        _store_id: rerun::external::re_log_types::StoreId,
+        path: std::path::PathBuf,
+        tx: std::sync::mpsc::Sender<re_data_source::LoadedData>,
+    ) -> Result<(), re_data_source::DataLoaderError> {
+        let contents = std::fs::read(&path)?;
+        if path.is_dir() {
+            return Err(re_data_source::DataLoaderError::Incompatible(path)); // simply not interested
+        }
+        hash_and_log(&tx, &path, &contents)
+    }
+
+    fn load_from_file_contents(
+        &self,
+        _store_id: rerun::external::re_log_types::StoreId,
+        filepath: std::path::PathBuf,
+        contents: std::borrow::Cow<'_, [u8]>,
+        tx: std::sync::mpsc::Sender<re_data_source::LoadedData>,
+    ) -> Result<(), re_data_source::DataLoaderError> {
+        hash_and_log(&tx, &filepath, &contents)
+    }
+}
+
+fn hash_and_log(
+    tx: &std::sync::mpsc::Sender<re_data_source::LoadedData>,
+    filepath: &std::path::Path,
+    contents: &[u8],
+) -> Result<(), re_data_source::DataLoaderError> {
+    use std::collections::hash_map::DefaultHasher;
+    use std::hash::{Hash, Hasher};
+
+    let mut h = DefaultHasher::new();
+    contents.hash(&mut h);
+
+    let doc = rerun::TextDocument::new(format!("{:08X}", h.finish()))
+        .with_media_type(rerun::MediaType::TEXT);
+
+    let entity_path = EntityPath::from_file_path(filepath);
+    let entity_path = format!("{entity_path}/hashed").into();
+    let row = DataRow::from_archetype(RowId::new(), TimePoint::timeless(), entity_path, &doc)?;
+
+    tx.send(row.into()).ok();
+
+    Ok(())
+}