Restructure python codegen: move codegen to rerun and old log APIs to log_deprecated

crates/re_types/build.rs

@@ -14,7 +14,8 @@ const ENTRYPOINT_PATH: &str = "./definitions/rerun/archetypes.fbs";
 const DOC_EXAMPLES_DIR_PATH: &str = "../../docs/code-examples";
 const CPP_OUTPUT_DIR_PATH: &str = "../../rerun_cpp";
 const RUST_OUTPUT_DIR_PATH: &str = ".";
-const PYTHON_OUTPUT_DIR_PATH: &str = "../../rerun_py/rerun_sdk/rerun/_rerun2";
+const PYTHON_OUTPUT_DIR_PATH: &str = "../../rerun_py/rerun_sdk/rerun";
+const PYTHON_TESTING_OUTPUT_DIR_PATH: &str = "../../rerun_py/tests/test_types";
 
 fn main() {
     if cfg!(target_os = "windows") {
@@ -92,6 +93,7 @@ fn main() {
         || {
             re_types_builder::generate_python_code(
                 PYTHON_OUTPUT_DIR_PATH,
+                PYTHON_TESTING_OUTPUT_DIR_PATH,
                 &objects,
                 &arrow_registry,
             );

crates/re_types_builder/src/bin/build_re_types.rs

@@ -6,7 +6,8 @@ const DEFINITIONS_DIR_PATH: &str = "crates/re_types/definitions";
 const ENTRYPOINT_PATH: &str = "crates/re_types/definitions/rerun/archetypes.fbs";
 const CPP_OUTPUT_DIR_PATH: &str = "rerun_cpp";
 const RUST_OUTPUT_DIR_PATH: &str = "crates/re_types/.";
-const PYTHON_OUTPUT_DIR_PATH: &str = "rerun_py/rerun_sdk/rerun/_rerun2";
+const PYTHON_OUTPUT_DIR_PATH: &str = "rerun_py/rerun_sdk/rerun";
+const PYTHON_TESTING_OUTPUT_DIR_PATH: &str = "rerun_py/tests/test_types";
 
 fn main() {
     re_log::setup_native_logging();
@@ -42,6 +43,7 @@ fn main() {
     let cpp_output_dir_path = workspace_dir.join(CPP_OUTPUT_DIR_PATH);
     let rust_output_dir_path = workspace_dir.join(RUST_OUTPUT_DIR_PATH);
     let python_output_dir_path = workspace_dir.join(PYTHON_OUTPUT_DIR_PATH);
+    let python_testing_output_dir_path = workspace_dir.join(PYTHON_TESTING_OUTPUT_DIR_PATH);
 
     re_log::info!("Running codegen…");
     let (objects, arrow_registry) =
@@ -54,6 +56,7 @@ fn main() {
         || {
             re_types_builder::generate_python_code(
                 python_output_dir_path,
+                python_testing_output_dir_path,
                 &objects,
                 &arrow_registry,
             );

crates/re_types_builder/src/codegen/python.rs

@@ -66,12 +66,14 @@ impl PythonObjectExt for Object {
 
 pub struct PythonCodeGenerator {
     pkg_path: Utf8PathBuf,
+    testing_pkg_path: Utf8PathBuf,
 }
 
 impl PythonCodeGenerator {
-    pub fn new(pkg_path: impl Into<Utf8PathBuf>) -> Self {
+    pub fn new(pkg_path: impl Into<Utf8PathBuf>, testing_pkg_path: impl Into<Utf8PathBuf>) -> Self {
         Self {
             pkg_path: pkg_path.into(),
+            testing_pkg_path: testing_pkg_path.into(),
         }
     }
 }
@@ -89,6 +91,8 @@ impl CodeGenerator for PythonCodeGenerator {
         }
 
         {
+            // TODO(jleibs): Should we still be generating an equivalent to this?
+            /*
             let archetype_names = objects
                 .ordered_objects(ObjectKind::Archetype.into())
                 .iter()
@@ -98,9 +102,10 @@ impl CodeGenerator for PythonCodeGenerator {
                 self.pkg_path.join("__init__.py"),
                 lib_source_code(&archetype_names),
             );
+            */
         }
 
-        write_files(&self.pkg_path, &files_to_write);
+        self.write_files(&files_to_write);
 
         let filepaths = files_to_write.keys().cloned().collect();
 
@@ -235,14 +240,20 @@ impl PythonCodeGenerator {
         files_to_write: &mut BTreeMap<Utf8PathBuf, String>,
     ) {
         let kind_path = self.pkg_path.join(object_kind.plural_snake_case());
+        let test_kind_path = self.testing_pkg_path.join(object_kind.plural_snake_case());
 
         // (module_name, [object_name])
         let mut mods = HashMap::<String, Vec<String>>::new();
+        let mut test_mods = HashMap::<String, Vec<String>>::new();
 
         // Generate folder contents:
         let ordered_objects = objects.ordered_objects(object_kind.into());
         for &obj in &ordered_objects {
-            let filepath = kind_path.join(format!("{}.py", obj.snake_case_name()));
+            let filepath = if obj.is_testing() {
+                test_kind_path.join(format!("{}.py", obj.snake_case_name()))
+            } else {
+                kind_path.join(format!("{}.py", obj.snake_case_name()))
+            };
 
             let ext_class = ExtensionClass::new(&kind_path, obj);
 
@@ -267,9 +278,14 @@ impl PythonCodeGenerator {
 
             // NOTE: Isolating the file stem only works because we're handling datatypes, components
             // and archetypes separately (and even then it's a bit shady, eh).
-            mods.entry(filepath.file_stem().unwrap().to_owned())
-                .or_default()
-                .extend(names.iter().cloned());
+            if obj.is_testing() {
+                &mut test_mods
+            } else {
+                &mut mods
+            }
+            .entry(filepath.file_stem().unwrap().to_owned())
+            .or_default()
+            .extend(names.iter().cloned());
 
             let mut code = String::new();
             code.push_text(&format!("# {}", autogen_warning!()), 1, 0);
@@ -287,8 +303,11 @@ impl PythonCodeGenerator {
 
             let manifest = quote_manifest(names);
 
+            let rerun_path = if obj.is_testing() { "rerun." } else { ".." };
+
             code.push_unindented_text(
-                "
+                format!(
+                    "
             from __future__ import annotations
 
             from typing import (Any, Dict, Iterable, Optional, Sequence, Set, Tuple, Union,
@@ -300,14 +319,14 @@ impl PythonCodeGenerator {
             import pyarrow as pa
             import uuid
 
-            from .._baseclasses import (
+            from {rerun_path}_baseclasses import (
                 Archetype,
                 BaseExtensionType,
                 BaseExtensionArray,
                 BaseDelegatingExtensionType,
                 BaseDelegatingExtensionArray
             )
-            from .._converters import (
+            from {rerun_path}_converters import (
                 int_or_none,
                 float_or_none,
                 bool_or_none,
@@ -325,7 +344,8 @@ impl PythonCodeGenerator {
                 to_np_float32,
                 to_np_float64
             )
-            ",
+            "
+                ),
                 0,
             );
 
@@ -366,60 +386,80 @@ impl PythonCodeGenerator {
         }
 
         // rerun/{datatypes|components|archetypes}/__init__.py
-        {
-            let path = kind_path.join("__init__.py");
+        write_init(&kind_path, &mods, files_to_write);
+        write_init(&test_kind_path, &test_mods, files_to_write);
+    }
 
-            let mut code = String::new();
+    fn write_files(&self, files_to_write: &BTreeMap<Utf8PathBuf, String>) {
+        re_tracing::profile_function!();
 
-            let manifest = quote_manifest(mods.iter().flat_map(|(_, names)| names.iter()));
+        // Running `black` once for each file is very slow, so we write all
+        // files to a temporary folder, format it, and copy back the results.
 
-            code.push_text(&format!("# {}", autogen_warning!()), 2, 0);
-            code.push_unindented_text(
-                "
-            from __future__ import annotations
+        let tempdir = tempfile::tempdir().unwrap();
+        let tempdir_path = Utf8PathBuf::try_from(tempdir.path().to_owned()).unwrap();
 
-            ",
-                0,
-            );
+        files_to_write.par_iter().for_each(|(filepath, source)| {
+            let formatted_source_path = self.format_path_for_tmp_dir(filepath, &tempdir_path);
+            super::common::write_file(&formatted_source_path, source);
+        });
 
-            for (module, names) in &mods {
-                let names = names.join(", ");
-                code.push_text(&format!("from .{module} import {names}"), 1, 0);
-            }
+        format_python_dir(&tempdir_path).unwrap();
 
-            code.push_unindented_text(format!("\n__all__ = [{manifest}]"), 0);
+        // Read back and copy to the final destination:
+        files_to_write
+            .par_iter()
+            .for_each(|(filepath, _original_source)| {
+                let formatted_source_path = self.format_path_for_tmp_dir(filepath, &tempdir_path);
+                let formatted_source = std::fs::read_to_string(formatted_source_path).unwrap();
+                super::common::write_file(filepath, &formatted_source);
+            });
+    }
 
-            files_to_write.insert(path, code);
-        }
+    fn format_path_for_tmp_dir(
+        &self,
+        filepath: &Utf8Path,
+        tempdir_path: &Utf8PathBuf,
+    ) -> Utf8PathBuf {
+        // If the prefix is pkg_path, strip it, and then append to tempdir
+        // However, if the prefix is testing_pkg_path, strip it and insert an extra
+        // "testing" to avoid name collisions.
+        filepath.strip_prefix(&self.pkg_path).map_or_else(
+            |_| {
+                tempdir_path
+                    .join("testing")
+                    .join(filepath.strip_prefix(&self.testing_pkg_path).unwrap())
+            },
+            |f| tempdir_path.join(f),
+        )
     }
 }
 
-fn write_files(pkg_path: &Utf8Path, files_to_write: &BTreeMap<Utf8PathBuf, String>) {
-    re_tracing::profile_function!();
-
-    // Running `black` once for each file is very slow, so we write all
-    // files to a temporary folder, format it, and copy back the results.
-
-    let tempdir = tempfile::tempdir().unwrap();
-    let tempdir_path = Utf8PathBuf::try_from(tempdir.path().to_owned()).unwrap();
-
-    files_to_write.par_iter().for_each(|(filepath, source)| {
-        let formatted_source_path = tempdir_path.join(filepath.strip_prefix(pkg_path).unwrap());
-        super::common::write_file(&formatted_source_path, source);
-    });
-
-    format_python_dir(&tempdir_path).unwrap();
+fn write_init(
+    kind_path: &Utf8PathBuf,
+    mods: &HashMap<String, Vec<String>>,
+    files_to_write: &mut BTreeMap<Utf8PathBuf, String>,
+) {
+    let path = kind_path.join("__init__.py");
+    let mut code = String::new();
+    let manifest = quote_manifest(mods.iter().flat_map(|(_, names)| names.iter()));
+    code.push_text(&format!("# {}", autogen_warning!()), 2, 0);
+    code.push_unindented_text(
+        "
+            from __future__ import annotations
 
-    // Read back and copy to the final destination:
-    files_to_write
-        .par_iter()
-        .for_each(|(filepath, _original_source)| {
-            let formatted_source_path = tempdir_path.join(filepath.strip_prefix(pkg_path).unwrap());
-            let formatted_source = std::fs::read_to_string(formatted_source_path).unwrap();
-            super::common::write_file(filepath, &formatted_source);
-        });
+            ",
+        0,
+    );
+    for (module, names) in mods {
+        let names = names.join(", ");
+        code.push_text(&format!("from .{module} import {names}"), 1, 0);
+    }
+    code.push_unindented_text(format!("\n__all__ = [{manifest}]"), 0);
+    files_to_write.insert(path, code);
 }
 
+#[allow(dead_code)]
 fn lib_source_code(archetype_names: &[String]) -> String {
     let manifest = quote_manifest(archetype_names);
     let archetype_names = archetype_names.join(", ");

crates/re_types_builder/src/lib.rs

@@ -260,7 +260,7 @@ pub fn generate_lang_agnostic(
 
 /// Generates a .gitattributes file that marks up all generated files as generated
 pub fn generate_gitattributes_for_generated_files(
-    output_path: impl AsRef<Utf8Path>,
+    output_path: &impl AsRef<Utf8Path>,
     files: impl Iterator<Item = Utf8PathBuf>,
 ) {
     let filename = ".gitattributes";
@@ -311,7 +311,7 @@ pub fn generate_cpp_code(
     re_tracing::profile_function!();
     let mut gen = CppCodeGenerator::new(output_path.as_ref());
     let filepaths = gen.generate(objects, arrow_registry);
-    generate_gitattributes_for_generated_files(output_path, filepaths.into_iter());
+    generate_gitattributes_for_generated_files(&output_path, filepaths.into_iter());
 }
 
 /// Generates Rust code.
@@ -340,7 +340,7 @@ pub fn generate_rust_code(
     re_tracing::profile_function!();
     let mut gen = RustCodeGenerator::new(output_crate_path.as_ref());
     let filepaths = gen.generate(objects, arrow_registry);
-    generate_gitattributes_for_generated_files(output_crate_path, filepaths.into_iter());
+    generate_gitattributes_for_generated_files(&output_crate_path, filepaths.into_iter());
 }
 
 /// Generates Python code.
@@ -356,20 +356,35 @@ pub fn generate_rust_code(
 ///     "./definitions/rerun/archetypes.fbs",
 /// );
 /// re_types_builder::generate_python_code(
-///     "./rerun_py",
+///     "./rerun_py/rerun_sdk",
+///     "./rerun_py/tests",
 ///     &objects,
 ///     &arrow_registry,
 /// );
 /// ```
 pub fn generate_python_code(
     output_pkg_path: impl AsRef<Utf8Path>,
+    testing_output_pkg_path: impl AsRef<Utf8Path>,
     objects: &Objects,
     arrow_registry: &ArrowRegistry,
 ) {
     re_tracing::profile_function!();
-    let mut gen = PythonCodeGenerator::new(output_pkg_path.as_ref());
+    let mut gen =
+        PythonCodeGenerator::new(output_pkg_path.as_ref(), testing_output_pkg_path.as_ref());
     let filepaths = gen.generate(objects, arrow_registry);
-    generate_gitattributes_for_generated_files(output_pkg_path, filepaths.into_iter());
+    generate_gitattributes_for_generated_files(
+        &output_pkg_path,
+        filepaths
+            .iter()
+            .filter(|f| f.starts_with(output_pkg_path.as_ref()))
+            .cloned(),
+    );
+    generate_gitattributes_for_generated_files(
+        &testing_output_pkg_path,
+        filepaths
+            .into_iter()
+            .filter(|f| f.starts_with(testing_output_pkg_path.as_ref())),
+    );
 }
 
 pub(crate) fn rerun_workspace_path() -> camino::Utf8PathBuf {

examples/python/depth_guided_stable_diffusion/huggingface_pipeline.py

@@ -57,7 +57,7 @@
 from transformers import DPTForDepthEstimation
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-logger.addHandler(rr.log.text.LoggingHandler("logs"))
+logger.addHandler(rr.LoggingHandler("logs"))
 
 
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.preprocess