From 7b2210832bf53a9645473620f7cdd966dd2a490a Mon Sep 17 00:00:00 2001
From: Aaron <aaron@skydio.com>
Date: Wed, 8 Mar 2023 17:32:43 +0000
Subject: [PATCH] [SymForce] Add Unit3 to cc_sym

Fixes https://github.com/symforce-org/symforce/issues/305

Topic: sf-unit3-py
Reviewers: bradley,nathan,chao,hayk
GitOrigin-RevId: b984e65be49ed8bbc14a1fbbdaa311d6de948844
---
 symforce/pybind/cc_sym.pyi                 | 9 +++++++++
 symforce/pybind/cc_values.cc               | 2 ++
 symforce/pybind/sym_type_casters.cc        | 1 +
 symforce/pybind/sym_type_casters.h         | 8 ++++++++
 test/symforce_cc_sym_stubs_codegen_test.py | 1 +
 test/symforce_values_test.py               | 1 +
 6 files changed, 22 insertions(+)

diff --git a/symforce/pybind/cc_sym.pyi b/symforce/pybind/cc_sym.pyi
index b9900787a..3c01a75f8 100644
--- 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
index ffbf3cabe..3f4e681fd 100644
--- 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
index 76152d4ec..63aac03b2 100644
--- 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
index 5cf468818..6928edeeb 100644
--- 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
index d61ca2151..35be2545d 100644
--- 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
index ce1fa5520..ba68cfbe9 100644
--- 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,