Support pickling geometry types

subprojects/robotpy-wpimath/gen/geometry/Ellipse2d.yml

@@ -42,7 +42,23 @@ inline_code: |
     .def_property_readonly("ysemiaxis_feet", [](Ellipse2d &self) -> units::foot_t {
       return self.YSemiAxis();
     })
-    .def("__repr__", py::overload_cast<const Ellipse2d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Ellipse2d&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Ellipse2d &self) { // __getstate__
+            return py::make_tuple(self.Center(), self.XSemiAxis(), self.YSemiAxis());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 3)
+                throw std::runtime_error("Invalid state!");
+
+            return Ellipse2d(
+              t[0].cast<Pose2d>(),
+              t[1].cast<units::meter_t>(),
+              t[2].cast<units::meter_t>()
+            );
+        }
+    ))
+    ;
 
 
   SetupWPyStruct<frc::Ellipse2d>(cls_Ellipse2d);

subprojects/robotpy-wpimath/gen/geometry/Pose2d.yml

@@ -62,6 +62,22 @@ inline_code: |
     .def_property_readonly("y_feet", [](Pose2d * self) -> units::foot_t {
       return self->Y();
     })
-    .def("__repr__", py::overload_cast<const Pose2d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Pose2d&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Pose2d &self) { // __getstate__
+            return py::make_tuple(self.X(), self.Y(), self.Rotation().Radians());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 3)
+                throw std::runtime_error("Invalid state!");
+
+            return Pose2d(
+              t[0].cast<units::meter_t>(),
+              t[1].cast<units::meter_t>(),
+              t[2].cast<units::radian_t>()
+            );
+        }
+    ))
+  ;
 
   SetupWPyStruct<frc::Pose2d>(cls_Pose2d);

subprojects/robotpy-wpimath/gen/geometry/Pose3d.yml

@@ -53,6 +53,31 @@ inline_code: |
     .def_property_readonly("z_feet", [](const Pose3d * self) -> units::foot_t {
       return self->Z();
     })
-    .def("__repr__", py::overload_cast<const Pose3d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Pose3d&>(&rpy::toString))
+    .def(py::pickle(
+      [](const Pose3d &self) { // __getstate__
+        auto q = self.Rotation().GetQuaternion();
+        return py::make_tuple(self.X(), self.Y(), self.Z(), q.W(), q.X(), q.Y(), q.Z());
+      },
+      [](py::tuple t) { // __setstate__
+        if (t.size() != 7)
+          throw std::runtime_error("Invalid state!");
+
+        return Pose3d(
+          t[0].cast<units::meter_t>(),
+          t[1].cast<units::meter_t>(),
+          t[2].cast<units::meter_t>(),
+          Rotation3d(
+            Quaternion(
+              t[3].cast<double>(),
+              t[4].cast<double>(),
+              t[5].cast<double>(),
+              t[6].cast<double>()
+            )
+          )
+        );
+      }
+    ))
+  ;
 
   SetupWPyStruct<frc::Pose3d>(cls_Pose3d);

subprojects/robotpy-wpimath/gen/geometry/Quaternion.yml

@@ -48,6 +48,23 @@ classes:
 
 inline_code: |
   cls_Quaternion
-    .def("__repr__", py::overload_cast<const Quaternion&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Quaternion&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Quaternion &self) { // __getstate__
+            return py::make_tuple(self.W(), self.X(), self.Y(), self.Z());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 4)
+                throw std::runtime_error("Invalid state!");
+
+            return Quaternion(
+              t[0].cast<double>(),
+              t[1].cast<double>(),
+              t[2].cast<double>(),
+              t[3].cast<double>()
+            );
+        }
+    ))
+  ;
 
   SetupWPyStruct<frc::Quaternion>(cls_Quaternion);

subprojects/robotpy-wpimath/gen/geometry/Rectangle2d.yml

@@ -41,7 +41,23 @@ inline_code: |
     .def_property_readonly("ywidth_feet", [](Rectangle2d &self) -> units::foot_t {
       return self.YWidth();
     })
-    .def("__repr__", py::overload_cast<const Rectangle2d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Rectangle2d&>(&rpy::toString))
+    .def(py::pickle(
+      [](const Rectangle2d &self) { // __getstate__
+        return py::make_tuple(self.Center(), self.XWidth(), self.YWidth());
+      },
+      [](py::tuple t) { // __setstate__
+        if (t.size() != 3)
+          throw std::runtime_error("Invalid state!");
+
+        return Rectangle2d(
+          t[0].cast<Pose2d>(),
+          t[1].cast<units::meter_t>(),
+          t[2].cast<units::meter_t>()
+        );
+      }
+    ))
+    ;
 
 
   SetupWPyStruct<frc::Rectangle2d>(cls_Rectangle2d);

subprojects/robotpy-wpimath/gen/geometry/Rotation2d.yml

@@ -51,6 +51,15 @@ inline_code: |
     .def_static("fromRotations", [](units::turn_t value) {
       return std::make_unique<Rotation2d>(value);
     })
-    .def("__repr__", py::overload_cast<const Rotation2d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Rotation2d&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Rotation2d &self) { // __getstate__
+            return self.Radians();
+        },
+        [](units::radian_t t) { // __setstate__
+            return Rotation2d(t);
+        }
+    ))
+  ;
 
   SetupWPyStruct<frc::Rotation2d>(cls_Rotation2d);

subprojects/robotpy-wpimath/gen/geometry/Rotation3d.yml

@@ -60,6 +60,15 @@ inline_code: |
     .def_property_readonly("angle_degrees", [](const Rotation3d * self) -> units::degree_t {
       return self->Angle();
     })
-    .def("__repr__", py::overload_cast<const Rotation3d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Rotation3d&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Rotation3d &self) { // __getstate__
+            return self.GetQuaternion();
+        },
+        [](Quaternion t) { // __setstate__
+            return Rotation3d(t);
+        }
+    ))
+  ;
 
   SetupWPyStruct<frc::Rotation3d>(cls_Rotation3d);

subprojects/robotpy-wpimath/gen/geometry/Transform2d.yml

@@ -46,7 +46,22 @@ inline_code: |
     .def_property_readonly("y_feet", [](Transform2d * self) -> units::foot_t {
       return self->Y();
     })
-    .def("__repr__", py::overload_cast<const Transform2d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Transform2d&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Transform2d &self) { // __getstate__
+            return py::make_tuple(self.X(), self.Y(), self.Rotation().Radians());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 3)
+                throw std::runtime_error("Invalid state!");
+
+            return Transform2d(
+              t[0].cast<units::meter_t>(),
+              t[1].cast<units::meter_t>(),
+              t[2].cast<units::radian_t>()
+            );
+        }
+    ))
   ;
 
   SetupWPyStruct<frc::Transform2d>(cls_Transform2d);

subprojects/robotpy-wpimath/gen/geometry/Transform3d.yml

@@ -45,6 +45,31 @@ inline_code: |
     .def_property_readonly("z_feet", [](const Transform3d * self) -> units::foot_t {
       return self->Z();
     })
-    .def("__repr__", py::overload_cast<const Transform3d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Transform3d&>(&rpy::toString))
+    .def(py::pickle(
+      [](const Transform3d &self) { // __getstate__
+        auto q = self.Rotation().GetQuaternion();
+        return py::make_tuple(self.X(), self.Y(), self.Z(), q.W(), q.X(), q.Y(), q.Z());
+      },
+      [](py::tuple t) { // __setstate__
+        if (t.size() != 7)
+          throw std::runtime_error("Invalid state!");
+
+        return Transform3d(
+          t[0].cast<units::meter_t>(),
+          t[1].cast<units::meter_t>(),
+          t[2].cast<units::meter_t>(),
+          Rotation3d(
+            Quaternion(
+              t[3].cast<double>(),
+              t[4].cast<double>(),
+              t[5].cast<double>(),
+              t[6].cast<double>()
+            )
+          )
+        );
+      }
+    ))
+  ;
 
   SetupWPyStruct<frc::Transform3d>(cls_Transform3d);

subprojects/robotpy-wpimath/gen/geometry/Translation2d.yml

@@ -80,6 +80,21 @@ inline_code: |
           throw std::out_of_range("Translation2d index out of range");
       }
     })
-    .def("__repr__", py::overload_cast<const Translation2d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Translation2d&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Translation2d &self) { // __getstate__
+            return py::make_tuple(self.X(), self.Y());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 2)
+                throw std::runtime_error("Invalid state!");
+
+            return Translation2d(
+              t[0].cast<units::meter_t>(),
+              t[1].cast<units::meter_t>()
+            );
+        }
+    ))
+  ;
 
   SetupWPyStruct<frc::Translation2d>(cls_Translation2d);

subprojects/robotpy-wpimath/gen/geometry/Translation3d.yml

@@ -74,6 +74,22 @@ inline_code: |
           throw std::out_of_range("Translation3d index out of range");
       }
     })
-    .def("__repr__", py::overload_cast<const Translation3d&>(&rpy::toString));
+    .def("__repr__", py::overload_cast<const Translation3d&>(&rpy::toString))
+    .def(py::pickle(
+        [](const Translation3d &self) { // __getstate__
+            return py::make_tuple(self.X(), self.Y(), self.Z());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 3)
+                throw std::runtime_error("Invalid state!");
+
+            return Translation3d(
+              t[0].cast<units::meter_t>(),
+              t[1].cast<units::meter_t>(),
+              t[2].cast<units::meter_t>()
+            );
+        }
+    ))
+  ;
 
   SetupWPyStruct<frc::Translation3d>(cls_Translation3d);

subprojects/robotpy-wpimath/gen/geometry/Twist2d.yml

@@ -55,6 +55,21 @@ inline_code: |
                      "dy=" + std::to_string(tw.dy()) + ", "
                      "dtheta=" + std::to_string(tw.dtheta()) + ")";
     })
+    .def(py::pickle(
+        [](const Twist2d &tw) { // __getstate__
+            return py::make_tuple(tw.dx(), tw.dy(), tw.dtheta());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 3)
+                throw std::runtime_error("Invalid state!");
+
+            return Twist2d(
+              t[0].cast<units::meter_t>(),
+              t[1].cast<units::meter_t>(),
+              t[2].cast<units::radian_t>()
+            );
+        }
+    ))
   ;
 
   SetupWPyStruct<frc::Twist2d>(cls_Twist2d);

subprojects/robotpy-wpimath/gen/geometry/Twist3d.yml

@@ -89,6 +89,24 @@ inline_code: |
                      "ry=" + std::to_string(tw.ry()) + ", "
                      "rz=" + std::to_string(tw.rz()) + ")";
     })
+    .def(py::pickle(
+        [](const Twist3d &tw) { // __getstate__
+            return py::make_tuple(tw.dx(), tw.dy(), tw.dz(), tw.rx(), tw.ry(), tw.rz());
+        },
+        [](py::tuple t) { // __setstate__
+            if (t.size() != 6)
+                throw std::runtime_error("Invalid state!");
+
+            return Twist3d(
+              t[0].cast<units::meter_t>(),
+              t[1].cast<units::meter_t>(),
+              t[2].cast<units::meter_t>(),
+              t[3].cast<units::radian_t>(),
+              t[4].cast<units::radian_t>(),
+              t[5].cast<units::radian_t>()
+            );
+        }
+    ))
   ;
 
   SetupWPyStruct<frc::Twist3d>(cls_Twist3d);