Skip to content

Commit 965f82d

Browse files
mailmindlinmailmindlin
committed
Support pickling geometry types
1 parent 75ded9d commit 965f82d

File tree

13 files changed

+223
-11
lines changed

13 files changed

+223
-11
lines changed

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,23 @@ inline_code: |
4242
.def_property_readonly("ysemiaxis_feet", [](Ellipse2d &self) -> units::foot_t {
4343
return self.YSemiAxis();
4444
})
45-
.def("__repr__", py::overload_cast<const Ellipse2d&>(&rpy::toString));
45+
.def("__repr__", py::overload_cast<const Ellipse2d&>(&rpy::toString))
46+
.def(py::pickle(
47+
[](const Ellipse2d &self) { // __getstate__
48+
return py::make_tuple(self.Center(), self.XSemiAxis(), self.YSemiAxis());
49+
},
50+
[](py::tuple t) { // __setstate__
51+
if (t.size() != 3)
52+
throw std::runtime_error("Invalid state!");
53+
54+
return Ellipse2d(
55+
t[0].cast<Pose2d>(),
56+
t[1].cast<units::meter_t>(),
57+
t[2].cast<units::meter_t>()
58+
);
59+
}
60+
))
61+
;
4662
4763
4864
SetupWPyStruct<frc::Ellipse2d>(cls_Ellipse2d);

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,22 @@ inline_code: |
6262
.def_property_readonly("y_feet", [](Pose2d * self) -> units::foot_t {
6363
return self->Y();
6464
})
65-
.def("__repr__", py::overload_cast<const Pose2d&>(&rpy::toString));
65+
.def("__repr__", py::overload_cast<const Pose2d&>(&rpy::toString))
66+
.def(py::pickle(
67+
[](const Pose2d &self) { // __getstate__
68+
return py::make_tuple(self.X(), self.Y(), self.Rotation().Radians());
69+
},
70+
[](py::tuple t) { // __setstate__
71+
if (t.size() != 3)
72+
throw std::runtime_error("Invalid state!");
73+
74+
return Pose2d(
75+
t[0].cast<units::meter_t>(),
76+
t[1].cast<units::meter_t>(),
77+
t[2].cast<units::radian_t>()
78+
);
79+
}
80+
))
81+
;
6682
6783
SetupWPyStruct<frc::Pose2d>(cls_Pose2d);

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

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,31 @@ inline_code: |
5353
.def_property_readonly("z_feet", [](const Pose3d * self) -> units::foot_t {
5454
return self->Z();
5555
})
56-
.def("__repr__", py::overload_cast<const Pose3d&>(&rpy::toString));
56+
.def("__repr__", py::overload_cast<const Pose3d&>(&rpy::toString))
57+
.def(py::pickle(
58+
[](const Pose3d &self) { // __getstate__
59+
auto q = self.Rotation().GetQuaternion();
60+
return py::make_tuple(self.X(), self.Y(), self.Z(), q.W(), q.X(), q.Y(), q.Z());
61+
},
62+
[](py::tuple t) { // __setstate__
63+
if (t.size() != 7)
64+
throw std::runtime_error("Invalid state!");
65+
66+
return Pose3d(
67+
t[0].cast<units::meter_t>(),
68+
t[1].cast<units::meter_t>(),
69+
t[2].cast<units::meter_t>(),
70+
Rotation3d(
71+
Quaternion(
72+
t[3].cast<double>(),
73+
t[4].cast<double>(),
74+
t[5].cast<double>(),
75+
t[6].cast<double>()
76+
)
77+
)
78+
);
79+
}
80+
))
81+
;
5782
5883
SetupWPyStruct<frc::Pose3d>(cls_Pose3d);

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,23 @@ classes:
4848

4949
inline_code: |
5050
cls_Quaternion
51-
.def("__repr__", py::overload_cast<const Quaternion&>(&rpy::toString));
51+
.def("__repr__", py::overload_cast<const Quaternion&>(&rpy::toString))
52+
.def(py::pickle(
53+
[](const Quaternion &self) { // __getstate__
54+
return py::make_tuple(self.W(), self.X(), self.Y(), self.Z());
55+
},
56+
[](py::tuple t) { // __setstate__
57+
if (t.size() != 4)
58+
throw std::runtime_error("Invalid state!");
59+
60+
return Quaternion(
61+
t[0].cast<double>(),
62+
t[1].cast<double>(),
63+
t[2].cast<double>(),
64+
t[3].cast<double>()
65+
);
66+
}
67+
))
68+
;
5269
5370
SetupWPyStruct<frc::Quaternion>(cls_Quaternion);

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,23 @@ inline_code: |
4141
.def_property_readonly("ywidth_feet", [](Rectangle2d &self) -> units::foot_t {
4242
return self.YWidth();
4343
})
44-
.def("__repr__", py::overload_cast<const Rectangle2d&>(&rpy::toString));
44+
.def("__repr__", py::overload_cast<const Rectangle2d&>(&rpy::toString))
45+
.def(py::pickle(
46+
[](const Rectangle2d &self) { // __getstate__
47+
return py::make_tuple(self.Center(), self.XWidth(), self.YWidth());
48+
},
49+
[](py::tuple t) { // __setstate__
50+
if (t.size() != 3)
51+
throw std::runtime_error("Invalid state!");
52+
53+
return Rectangle2d(
54+
t[0].cast<Pose2d>(),
55+
t[1].cast<units::meter_t>(),
56+
t[2].cast<units::meter_t>()
57+
);
58+
}
59+
))
60+
;
4561
4662
4763
SetupWPyStruct<frc::Rectangle2d>(cls_Rectangle2d);

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,15 @@ inline_code: |
5151
.def_static("fromRotations", [](units::turn_t value) {
5252
return std::make_unique<Rotation2d>(value);
5353
})
54-
.def("__repr__", py::overload_cast<const Rotation2d&>(&rpy::toString));
54+
.def("__repr__", py::overload_cast<const Rotation2d&>(&rpy::toString))
55+
.def(py::pickle(
56+
[](const Rotation2d &self) { // __getstate__
57+
return self.Radians();
58+
},
59+
[](units::radian_t t) { // __setstate__
60+
return Rotation2d(t);
61+
}
62+
))
63+
;
5564
5665
SetupWPyStruct<frc::Rotation2d>(cls_Rotation2d);

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,15 @@ inline_code: |
6060
.def_property_readonly("angle_degrees", [](const Rotation3d * self) -> units::degree_t {
6161
return self->Angle();
6262
})
63-
.def("__repr__", py::overload_cast<const Rotation3d&>(&rpy::toString));
63+
.def("__repr__", py::overload_cast<const Rotation3d&>(&rpy::toString))
64+
.def(py::pickle(
65+
[](const Rotation3d &self) { // __getstate__
66+
return self.GetQuaternion();
67+
},
68+
[](Quaternion t) { // __setstate__
69+
return Rotation3d(t);
70+
}
71+
))
72+
;
6473
6574
SetupWPyStruct<frc::Rotation3d>(cls_Rotation3d);

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,22 @@ inline_code: |
4646
.def_property_readonly("y_feet", [](Transform2d * self) -> units::foot_t {
4747
return self->Y();
4848
})
49-
.def("__repr__", py::overload_cast<const Transform2d&>(&rpy::toString));
49+
.def("__repr__", py::overload_cast<const Transform2d&>(&rpy::toString))
50+
.def(py::pickle(
51+
[](const Transform2d &self) { // __getstate__
52+
return py::make_tuple(self.X(), self.Y(), self.Rotation().Radians());
53+
},
54+
[](py::tuple t) { // __setstate__
55+
if (t.size() != 3)
56+
throw std::runtime_error("Invalid state!");
57+
58+
return Transform2d(
59+
t[0].cast<units::meter_t>(),
60+
t[1].cast<units::meter_t>(),
61+
t[2].cast<units::radian_t>()
62+
);
63+
}
64+
))
5065
;
5166
5267
SetupWPyStruct<frc::Transform2d>(cls_Transform2d);

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

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,31 @@ inline_code: |
4545
.def_property_readonly("z_feet", [](const Transform3d * self) -> units::foot_t {
4646
return self->Z();
4747
})
48-
.def("__repr__", py::overload_cast<const Transform3d&>(&rpy::toString));
48+
.def("__repr__", py::overload_cast<const Transform3d&>(&rpy::toString))
49+
.def(py::pickle(
50+
[](const Transform3d &self) { // __getstate__
51+
auto q = self.Rotation().GetQuaternion();
52+
return py::make_tuple(self.X(), self.Y(), self.Z(), q.W(), q.X(), q.Y(), q.Z());
53+
},
54+
[](py::tuple t) { // __setstate__
55+
if (t.size() != 7)
56+
throw std::runtime_error("Invalid state!");
57+
58+
return Transform3d(
59+
t[0].cast<units::meter_t>(),
60+
t[1].cast<units::meter_t>(),
61+
t[2].cast<units::meter_t>(),
62+
Rotation3d(
63+
Quaternion(
64+
t[3].cast<double>(),
65+
t[4].cast<double>(),
66+
t[5].cast<double>(),
67+
t[6].cast<double>()
68+
)
69+
)
70+
);
71+
}
72+
))
73+
;
4974
5075
SetupWPyStruct<frc::Transform3d>(cls_Transform3d);

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

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,21 @@ inline_code: |
8080
throw std::out_of_range("Translation2d index out of range");
8181
}
8282
})
83-
.def("__repr__", py::overload_cast<const Translation2d&>(&rpy::toString));
83+
.def("__repr__", py::overload_cast<const Translation2d&>(&rpy::toString))
84+
.def(py::pickle(
85+
[](const Translation2d &self) { // __getstate__
86+
return py::make_tuple(self.X(), self.Y());
87+
},
88+
[](py::tuple t) { // __setstate__
89+
if (t.size() != 2)
90+
throw std::runtime_error("Invalid state!");
91+
92+
return Translation2d(
93+
t[0].cast<units::meter_t>(),
94+
t[1].cast<units::meter_t>()
95+
);
96+
}
97+
))
98+
;
8499
85100
SetupWPyStruct<frc::Translation2d>(cls_Translation2d);

0 commit comments

Comments
 (0)