Make TrajectoryConstraint.MinMax namedtuple-like

robotpy · Jan 1, 2025 · 65a26e5 · 65a26e5
1 parent 4313900
commit 65a26e5
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 14 deletions.
diff --git a/subprojects/robotpy-wpimath/gen/controls/TrajectoryConstraint.yml b/subprojects/robotpy-wpimath/gen/controls/TrajectoryConstraint.yml
@@ -15,3 +15,26 @@ classes:
     attributes:
       minAcceleration:
       maxAcceleration:
+    inline_code: |
+      .def(py::init([](
+          units::meters_per_second_squared_t minAcceleration,
+          units::meters_per_second_squared_t maxAcceleration) {
+        return frc::TrajectoryConstraint::MinMax{minAcceleration, maxAcceleration};
+      }), py::arg("minAcceleration"), py::arg("maxAcceleration"))
+
+      .def("__len__", [](const frc::TrajectoryConstraint::MinMax& self) { return 2; })
+      .def("__getitem__", [](const frc::TrajectoryConstraint::MinMax& self, int index) {
+        switch (index) {
+          case 0:
+            return self.minAcceleration;
+          case 1:
+            return self.maxAcceleration;
+          default:
+            throw std::out_of_range("TrajectoryConstraint.MinMax index out of range");
+        }
+      })
+
+      .def("__repr__", [](const frc::TrajectoryConstraint::MinMax &self) {
+        return py::str("TrajectoryConstraint.MinMax(minAcceleration={}, maxAcceleration={})").format(
+          self.minAcceleration, self.maxAcceleration);
+      })
diff --git a/subprojects/robotpy-wpimath/tests/test_trajectory.py b/subprojects/robotpy-wpimath/tests/test_trajectory.py
@@ -1,22 +1,27 @@
-from wpimath.trajectory import TrajectoryGenerator, TrajectoryParameterizer
-from wpimath.geometry import Pose2d, Transform2d
-import wpimath.trajectory
-import wpimath.trajectory.constraint
-from wpimath.spline import CubicHermiteSpline, SplineHelper
-
-from wpimath.geometry import Ellipse2d, Pose2d, Rectangle2d, Rotation2d, Translation2d
-
-from wpimath.trajectory import Trajectory, TrajectoryConfig
+import math
 
+from wpimath.geometry import (
+    Ellipse2d,
+    Pose2d,
+    Rectangle2d,
+    Rotation2d,
+    Transform2d,
+    Translation2d,
+)
+from wpimath.spline import CubicHermiteSpline, SplineHelper
+from wpimath.trajectory import (
+    Trajectory,
+    TrajectoryConfig,
+    TrajectoryGenerator,
+    TrajectoryParameterizer,
+)
 from wpimath.trajectory.constraint import (
-    MaxVelocityConstraint,
     EllipticalRegionConstraint,
+    MaxVelocityConstraint,
     RectangularRegionConstraint,
+    TrajectoryConstraint,
 )
 
-import math
-import typing
-
 
 def getTestTrajectory(config: TrajectoryConfig) -> Trajectory:
     # 2018 cross scale auto waypoints
@@ -100,6 +105,22 @@ def test_rectangular_region_constraint():
     assert exceededConstraintOutsideRegion
 
 
+#
+# TrajectoryConstraint
+#
+
+
+def test_trajectory_constraint_min_max():
+    min_max = TrajectoryConstraint.MinMax()
+    _, _ = min_max
+
+    min_max = TrajectoryConstraint.MinMax(minAcceleration=0, maxAcceleration=1)
+    assert (
+        repr(min_max)
+        == "TrajectoryConstraint.MinMax(minAcceleration=0.0, maxAcceleration=1.0)"
+    )
+
+
 #
 # TrajectoryParameterizer
 #
@@ -114,7 +135,7 @@ def test_trajectory_parameterizer():
     spline = CubicHermiteSpline(vec1.x, vec2.x, vec1.y, vec2.y)
 
     # sample the pose and curvature along the spline
-    points: typing.Tuple[Pose2d, float] = []
+    points: list[tuple[Pose2d, float]] = []
     for i in range(100):
         points.append(spline.getPoint(i / 100))