[SymForce] Add Unit3 to cc_sym

Fixes #305 Topic: sf-unit3-py Reviewers: bradley,nathan,chao,hayk GitOrigin-RevId: b984e65be49ed8bbc14a1fbbdaa311d6de948844
symforce-org · Mar 8, 2023 · 7b22108 · 7b22108
1 parent c1fdca1
commit 7b22108
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 0 deletions.
diff --git a/symforce/pybind/cc_sym.pyi b/symforce/pybind/cc_sym.pyi
@@ -26,6 +26,7 @@ from sym import Pose3
 from sym import Rot2
 from sym import Rot3
 from sym import SphericalCameraCal
+from sym import Unit3
 
 # isort: off
 
@@ -956,6 +957,10 @@ class Values:
 
         Add or update a value by key. Returns true if added, false if updated.
 
+        Update a value by index entry with no map lookup (compared to Set(key)). This does NOT add new values and assumes the key exists already.
+
+        Add or update a value by key. Returns true if added, false if updated.
+
         Update a value by index entry with no map lookup (compared to Set(key)). This does NOT add new values and assumes the key exists already.
         """
     @typing.overload
@@ -977,6 +982,8 @@ class Values:
     @typing.overload
     def set(self, key: Key, value: SphericalCameraCal) -> bool: ...
     @typing.overload
+    def set(self, key: Key, value: Unit3) -> bool: ...
+    @typing.overload
     def set(self, key: Key, value: float) -> bool: ...
     @typing.overload
     def set(self, key: Key, value: numpy.ndarray) -> bool: ...
@@ -1001,6 +1008,8 @@ class Values:
     @typing.overload
     def set(self, key: index_entry_t, value: SphericalCameraCal) -> None: ...
     @typing.overload
+    def set(self, key: index_entry_t, value: Unit3) -> None: ...
+    @typing.overload
     def set(self, key: index_entry_t, value: float) -> None: ...
     @typing.overload
     def set(self, key: index_entry_t, value: numpy.ndarray) -> None: ...

diff --git a/symforce/pybind/cc_values.cc b/symforce/pybind/cc_values.cc
@@ -34,6 +34,7 @@
 #include <sym/rot2.h>
 #include <sym/rot3.h>
 #include <sym/spherical_camera_cal.h>
+#include <sym/unit3.h>
 #include <sym/util/type_ops.h>
 #include <symforce/opt/key.h>
 #include <symforce/opt/values.h>
@@ -274,6 +275,7 @@ void AddValuesWrapper(pybind11::module_ module) {
   RegisterTypeWithValues<sym::Rot3d>(values_class);
   RegisterTypeWithValues<sym::Pose2d>(values_class);
   RegisterTypeWithValues<sym::Pose3d>(values_class);
+  RegisterTypeWithValues<sym::Unit3d>(values_class);
   RegisterTypeWithValues<sym::ATANCameraCald>(values_class);
   RegisterTypeWithValues<sym::DoubleSphereCameraCald>(values_class);
   RegisterTypeWithValues<sym::EquirectangularCameraCald>(values_class);

diff --git a/symforce/pybind/sym_type_casters.cc b/symforce/pybind/sym_type_casters.cc
@@ -12,6 +12,7 @@ constexpr descr<4> handle_sym_type_name<sym::Rot2d>::name;
 constexpr descr<4> handle_sym_type_name<sym::Rot3d>::name;
 constexpr descr<5> handle_sym_type_name<sym::Pose2d>::name;
 constexpr descr<5> handle_sym_type_name<sym::Pose3d>::name;
+constexpr descr<5> handle_sym_type_name<sym::Unit3d>::name;
 constexpr descr<13> handle_sym_type_name<sym::ATANCameraCald>::name;
 constexpr descr<21> handle_sym_type_name<sym::DoubleSphereCameraCald>::name;
 constexpr descr<24> handle_sym_type_name<sym::EquirectangularCameraCald>::name;

diff --git a/symforce/pybind/sym_type_casters.h b/symforce/pybind/sym_type_casters.h
@@ -27,6 +27,7 @@
 #include <sym/rot2.h>
 #include <sym/rot3.h>
 #include <sym/spherical_camera_cal.h>
+#include <sym/unit3.h>
 
 namespace py = pybind11;
 
@@ -59,6 +60,11 @@ struct handle_sym_type_name<sym::Pose3d> {
   static constexpr auto name = _("Pose3");
 };
 
+template <>
+struct handle_sym_type_name<sym::Unit3d> {
+  static constexpr auto name = _("Unit3");
+};
+
 template <>
 struct handle_sym_type_name<sym::ATANCameraCald> {
   static constexpr auto name = _("ATANCameraCal");
@@ -133,6 +139,8 @@ struct type_caster<sym::Pose2d> : public sym_type_caster<sym::Pose2d> {};
 template <>
 struct type_caster<sym::Pose3d> : public sym_type_caster<sym::Pose3d> {};
 template <>
+struct type_caster<sym::Unit3d> : public sym_type_caster<sym::Unit3d> {};
+template <>
 struct type_caster<sym::ATANCameraCald> : public sym_type_caster<sym::ATANCameraCald> {};
 template <>
 struct type_caster<sym::DoubleSphereCameraCald>

diff --git a/test/symforce_cc_sym_stubs_codegen_test.py b/test/symforce_cc_sym_stubs_codegen_test.py
@@ -103,6 +103,7 @@ class TypeStubParts:
                         "Rot2",
                         "Rot3",
                         "SphericalCameraCal",
+                        "Unit3",
                     ],
                     third_party_includes=["import scipy"],
                     cleaned_up_stubgen_output=stubgen_output,

diff --git a/test/symforce_values_test.py b/test/symforce_values_test.py
@@ -551,6 +551,7 @@ def test_from_storage_index(self) -> None:
                 sf.Rot3,
                 sf.Pose2,
                 sf.Pose3,
+                sf.Unit3,
                 sf.Complex,
                 sf.Quaternion,
                 sf.DualQuaternion,