Python codegen: support for Transform3D (#2639)

### What

This PR adds support for `Transform3D` to the new-gen Python SDK,
including:
- adding the relevant fbs
- implementing codegen'd union types
- adding all required hand-coded overrides
- adding typed-checked test to demonstrate/assess API
- full typing compliance for both mypy and pyright (required
`disallow_untyped_calls = False` in `.mypy.ini`)

**Note**:
- The `from_parent` field is moved from `Transform3D` to both
`Translation*` objects, to comply with the 1-comp-1-datatype rule.
- The `from_parent` field has been made `nullable` though it has a
default value, which are not yet implemented (#2641).

TODO:
- [x] arrow serialisation (basic, to be tested)
- [x] `rr.arch.Transform3D()` serialisation tests
- [x] default types (`from_parent`) (including proper field sorting)
- [x] ~~`ruff --fix` seems to always fail~~ (config error on my side)
 

TODO @teh-cmc and/or @jleibs :
- [x] deal with commented-out hack in `rerun_py/src/arrow.rs`
- [x] fix Rust codegen
- [x] fix restriction on nullable union (see discussion in #2619)
- [x] teach me how to properly indent text (in `python.rs`)


### Checklist
* [x] I have read and agree to [Contributor
Guide](https://github.com/rerun-io/rerun/blob/main/CONTRIBUTING.md) and
the [Code of
Conduct](https://github.com/rerun-io/rerun/blob/main/CODE_OF_CONDUCT.md)
* [x] I've included a screenshot or gif (if applicable)
* [x] I have tested [demo.rerun.io](https://demo.rerun.io/pr/2639) (if
applicable)

- [PR Build Summary](https://build.rerun.io/pr/2639)
- [Docs
preview](https://rerun.io/preview/pr%3Aantoine%2Fhope-python-transforms-new/docs)
- [Examples
preview](https://rerun.io/preview/pr%3Aantoine%2Fhope-python-transforms-new/examples)

---------

Co-authored-by: Jeremy Leibs <jeremy@rerun.io>

.mypy.ini

@@ -8,11 +8,8 @@ enable_error_code = redundant-expr, truthy-bool, ignore-without-code
 plugins = numpy.typing.mypy_plugin
 ignore_missing_imports = True
 no_implicit_reexport = False
+disallow_untyped_calls = False
 
 # Don't lint demo helpers.
 [mypy-rerun_demo.*]
 ignore_errors = True
-
-# Allow DNA example to call demo helpers.
-[mypy-python.dna.main]
-disallow_untyped_calls = False

crates/re_types/build.rs

@@ -21,6 +21,7 @@ const DOC_EXAMPLES_DIR_PATH: &str = "../../docs/code-examples";
 const CPP_OUTPUT_DIR_PATH: &str = "../../rerun_cpp/src";
 const RUST_OUTPUT_DIR_PATH: &str = ".";
 const PYTHON_OUTPUT_DIR_PATH: &str = "../../rerun_py/rerun_sdk/rerun/_rerun2";
+const PYTHON_PYPROJECT_PATH: &str = "../../rerun_py/pyproject.toml";
 
 // located in PYTHON_OUTPUT_DIR_PATH
 const ARCHETYPE_OVERRIDES_SUB_DIR_PATH: &str = "archetypes/_overrides";
@@ -120,17 +121,6 @@ fn generate_and_format_python_code(
 ) {
     re_types_builder::generate_python_code(PYTHON_OUTPUT_DIR_PATH, objects, arrow_registry);
 
-    let pyproject_path = PathBuf::from(PYTHON_OUTPUT_DIR_PATH)
-        .parent()
-        .unwrap()
-        .parent()
-        .unwrap()
-        .parent()
-        .unwrap()
-        .join("pyproject.toml")
-        .to_string_lossy()
-        .to_string();
-
     // TODO(emilk): format the python code _before_ writing them to file instead,
     // just like we do with C++ and Rust.
     // This should be doable py piping the code of each file to black/ruff via stdin.
@@ -152,17 +142,17 @@ fn generate_and_format_python_code(
     // 3) Call black again to cleanup some whitespace issues ruff might introduce
 
     let sh = Shell::new().unwrap();
-    call_black(&sh, &pyproject_path);
-    call_ruff(&sh, &pyproject_path);
-    call_black(&sh, &pyproject_path);
+    call_black(&sh);
+    call_ruff(&sh);
+    call_black(&sh);
 }
 
 // Do 3 things in parallel
 fn join3(a: impl FnOnce() + Send, b: impl FnOnce() + Send, c: impl FnOnce() + Send) {
     rayon::join(a, || rayon::join(b, c));
 }
 
-fn call_black(sh: &Shell, pyproject_path: &String) {
+fn call_black(sh: &Shell) {
     // NOTE: We're purposefully ignoring the error here.
     //
     // If the user doesn't have `black` in their $PATH, there's still no good reason to fail
@@ -171,13 +161,13 @@ fn call_black(sh: &Shell, pyproject_path: &String) {
     // The CI will catch the unformatted files at PR time and complain appropriately anyhow.
     cmd!(
         sh,
-        "black --config {pyproject_path} {PYTHON_OUTPUT_DIR_PATH}"
+        "black --config {PYTHON_PYPROJECT_PATH} {PYTHON_OUTPUT_DIR_PATH}"
     )
     .run()
     .ok();
 }
 
-fn call_ruff(sh: &Shell, pyproject_path: &String) {
+fn call_ruff(sh: &Shell) {
     // NOTE: We're purposefully ignoring the error here.
     //
     // If the user doesn't have `ruff` in their $PATH, there's still no good reason to fail
@@ -186,7 +176,7 @@ fn call_ruff(sh: &Shell, pyproject_path: &String) {
     // The CI will catch the unformatted files at PR time and complain appropriately anyhow.
     cmd!(
         sh,
-        "ruff --config {pyproject_path} --fix {PYTHON_OUTPUT_DIR_PATH}"
+        "ruff --config {PYTHON_PYPROJECT_PATH} --fix {PYTHON_OUTPUT_DIR_PATH}"
     )
     .run()
     .ok();

crates/re_types/definitions/rerun/components/transform3d.fbs

@@ -8,6 +8,10 @@ namespace rerun.components;
 
 /// An affine transform between two 3D spaces, represented in a given direction.
 table Transform3D (
+  // TODO(jleibs) We cannot set the fqname to rerun.transform3d until we have aligned
+  // the arrow schemas. Otherwise the schema-patching in `array_to_rust` will set the
+  // wrong schema for this component.
+  //"attr.rerun.legacy_fqname": "rerun.transform3d",
   order: 100
 ) {
   /// Representation of the transform.

crates/re_types/definitions/rerun/datatypes/point2d.fbs

@@ -10,7 +10,7 @@ namespace rerun.datatypes;
 /// A point in 2D space.
 struct Point2D (
   "attr.python.aliases": "Sequence[float]",
-  "attr.python.array_aliases": "npt.NDArray[np.float32], Sequence[npt.NDArray[np.float32]], Sequence[Tuple[float, float]], Sequence[float]",
+  "attr.python.array_aliases": "npt.NDArray[Any], Sequence[npt.NDArray[Any]], Sequence[Tuple[float, float]], Sequence[float]",
   "attr.rust.derive": "Default, Copy, PartialEq, PartialOrd",
   order: 100
 ) {

crates/re_types/definitions/rerun/datatypes/rotation3d.fbs

@@ -9,6 +9,7 @@ namespace rerun.datatypes;
 /// A 3D rotation.
 union Rotation3D (
   "attr.rust.derive": "Copy, PartialEq",
+  "attr.python.aliases": "Sequence[SupportsFloat]",
   order: 400
 ) {
   /// Rotation defined by a quaternion.

crates/re_types/definitions/rerun/datatypes/translation_and_mat3x3.fbs

@@ -15,13 +15,19 @@ table TranslationAndMat3x3 (
   order: 500
 ) {
   /// 3D translation, applied after the matrix.
+  //
+  // NOTE: Nullable rather than defaulting to an identity-like value because we want to be able
+  // to differentiate between no value vs. default value in the backend.
   translation: rerun.datatypes.Vec3D (nullable, order: 100);
 
   /// 3x3 matrix for scale, rotation & shear.
+  //
+  // NOTE: Nullable rather than defaulting to an identity-like value because we want to be able
+  // to differentiate between no value vs. default value in the backend.
   matrix: Mat3x3 (nullable, order: 200);
 
   // TODO(#2641): make this field non-nullable when default values are supported
   /// If true, the transform maps from the parent space to the space where the transform was logged.
   /// Otherwise, the transform maps from the space to its parent.
-  from_parent: bool = false (nullable, order: 300);
+  from_parent: bool = false (order: 300);
 }

crates/re_types/definitions/rerun/datatypes/translation_rotation_scale3d.fbs

@@ -15,16 +15,25 @@ table TranslationRotationScale3D (
   order: 600
 ) {
   /// 3D translation vector, applied last.
+  //
+  // NOTE: Nullable rather than defaulting to an identity-like value because we want to be able
+  // to differentiate between no value vs. default value in the backend.
   translation: rerun.datatypes.Vec3D (nullable, order: 100);
 
   /// 3D rotation, applied second.
+  //
+  // NOTE: Nullable rather than defaulting to an identity-like value because we want to be able
+  // to differentiate between no value vs. default value in the backend.
   rotation: rerun.datatypes.Rotation3D (nullable, order: 200);
 
   /// 3D scale, applied first.
+  //
+  // NOTE: Nullable rather than defaulting to an identity-like value because we want to be able
+  // to differentiate between no value vs. default value in the backend.
   scale: rerun.datatypes.Scale3D (nullable, order: 300);
 
   // TODO(#2641): make this field non-nullable when default values are supported
   /// If true, the transform maps from the parent space to the space where the transform was logged.
   /// Otherwise, the transform maps from the space to its parent.
-  from_parent: bool = false (nullable, order: 400);
+  from_parent: bool = false (order: 400);
 }

crates/re_types/source_hash.txt

@@ -1,4 +1,4 @@
 # This is a sha256 hash for all direct and indirect dependencies of this crate's build script.
 # It can be safely removed at anytime to force the build script to run again.
 # Check out build.rs to see how it's computed.
-48b1929d5cb17125eaae7733df116017fed8e27f5202d2365146e68e9a1a5b16
+114ba22844997f4dc1a9fc81ad83419d80cf124c412777c44e1b9b71b3e50510

crates/re_types/src/components/transform3d.rs

@@ -80,7 +80,7 @@ impl crate::Loggable for Transform3D {
                         Field {
                             name: "from_parent".to_owned(),
                             data_type: DataType::Boolean,
-                            is_nullable: true,
+                            is_nullable: false,
                             metadata: [].into(),
                         },
                     ]),
@@ -207,7 +207,7 @@ impl crate::Loggable for Transform3D {
                         Field {
                             name: "from_parent".to_owned(),
                             data_type: DataType::Boolean,
-                            is_nullable: true,
+                            is_nullable: false,
                             metadata: [].into(),
                         },
                     ]),

crates/re_types/src/datatypes/transform3d.rs

@@ -83,7 +83,7 @@ impl crate::Loggable for Transform3D {
                         Field {
                             name: "from_parent".to_owned(),
                             data_type: DataType::Boolean,
-                            is_nullable: true,
+                            is_nullable: false,
                             metadata: [].into(),
                         },
                     ]),
@@ -210,7 +210,7 @@ impl crate::Loggable for Transform3D {
                         Field {
                             name: "from_parent".to_owned(),
                             data_type: DataType::Boolean,
-                            is_nullable: true,
+                            is_nullable: false,
                             metadata: [].into(),
                         },
                     ]),

crates/re_types/src/datatypes/translation_and_mat3x3.rs

@@ -25,7 +25,7 @@ pub struct TranslationAndMat3x3 {
 
     /// If true, the transform maps from the parent space to the space where the transform was logged.
     /// Otherwise, the transform maps from the space to its parent.
-    pub from_parent: Option<bool>,
+    pub from_parent: bool,
 }
 
 impl<'a> From<TranslationAndMat3x3> for ::std::borrow::Cow<'a, TranslationAndMat3x3> {
@@ -85,7 +85,7 @@ impl crate::Loggable for TranslationAndMat3x3 {
             Field {
                 name: "from_parent".to_owned(),
                 data_type: DataType::Boolean,
-                is_nullable: true,
+                is_nullable: false,
                 metadata: [].into(),
             },
         ])
@@ -258,13 +258,10 @@ impl crate::Loggable for TranslationAndMat3x3 {
                         let (somes, from_parent): (Vec<_>, Vec<_>) = data
                             .iter()
                             .map(|datum| {
-                                let datum = datum
-                                    .as_ref()
-                                    .map(|datum| {
-                                        let Self { from_parent, .. } = &**datum;
-                                        from_parent.clone()
-                                    })
-                                    .flatten();
+                                let datum = datum.as_ref().map(|datum| {
+                                    let Self { from_parent, .. } = &**datum;
+                                    from_parent.clone()
+                                });
                                 (datum.is_some(), datum)
                             })
                             .unzip();
@@ -384,7 +381,11 @@ impl crate::Loggable for TranslationAndMat3x3 {
                                 Ok(Self {
                                     translation,
                                     matrix,
-                                    from_parent,
+                                    from_parent: from_parent.ok_or_else(|| {
+                                        crate::DeserializationError::MissingData {
+                                            datatype: data.data_type().clone(),
+                                        }
+                                    })?,
                                 })
                             })
                             .transpose()

crates/re_types/src/datatypes/translation_rotation_scale3d.rs

@@ -26,7 +26,7 @@ pub struct TranslationRotationScale3D {
 
     /// If true, the transform maps from the parent space to the space where the transform was logged.
     /// Otherwise, the transform maps from the space to its parent.
-    pub from_parent: Option<bool>,
+    pub from_parent: bool,
 }
 
 impl<'a> From<TranslationRotationScale3D> for ::std::borrow::Cow<'a, TranslationRotationScale3D> {
@@ -174,7 +174,7 @@ impl crate::Loggable for TranslationRotationScale3D {
             Field {
                 name: "from_parent".to_owned(),
                 data_type: DataType::Boolean,
-                is_nullable: true,
+                is_nullable: false,
                 metadata: [].into(),
             },
         ])
@@ -333,13 +333,10 @@ impl crate::Loggable for TranslationRotationScale3D {
                         let (somes, from_parent): (Vec<_>, Vec<_>) = data
                             .iter()
                             .map(|datum| {
-                                let datum = datum
-                                    .as_ref()
-                                    .map(|datum| {
-                                        let Self { from_parent, .. } = &**datum;
-                                        from_parent.clone()
-                                    })
-                                    .flatten();
+                                let datum = datum.as_ref().map(|datum| {
+                                    let Self { from_parent, .. } = &**datum;
+                                    from_parent.clone()
+                                });
                                 (datum.is_some(), datum)
                             })
                             .unzip();
@@ -446,7 +443,11 @@ impl crate::Loggable for TranslationRotationScale3D {
                                     translation,
                                     rotation,
                                     scale,
-                                    from_parent,
+                                    from_parent: from_parent.ok_or_else(|| {
+                                        crate::DeserializationError::MissingData {
+                                            datatype: data.data_type().clone(),
+                                        }
+                                    })?,
                                 })
                             })
                             .transpose()