New shorthand to get bodies from a Arbiter. #279

Author: viblo

CHANGELOG.rst

@@ -18,6 +18,7 @@ Changelog
    Added default do_nothing and always_collide callback functions to the CollisionHandler, so that its clear how to reset and align with other callbacks.
    If in old code you did handler.begin = None, you should now instead to handler.begin = CollisionHandler.always_collide etc.
 
+   New feature: Arbiter.bodies shorthand to ge the shapes' bodies in the Arbiter
    Changed type of PointQueryInfo.shape, SegmentQueryInfo.shape and ShapeQueryInfo.shape to not be Optional, they will always have a shape.
    New feature: ShapeFilter.rejects_collision()
    New feature: Added Vec2d.polar_tuple

pymunk/arbiter.py

@@ -5,6 +5,7 @@
 
 if TYPE_CHECKING:
     from .space import Space
+    from .body import Body
 
 from ._chipmunk_cffi import ffi, lib
 from .contact_point_set import ContactPointSet
@@ -69,6 +70,24 @@ def contact_point_set(self, point_set: ContactPointSet) -> None:
 
         lib.cpArbiterSetContactPointSet(self._arbiter, ffi.addressof(_set))
 
+    @property
+    def bodies(self) -> Tuple["Body", "Body"]:
+        """The the bodies in the order their corresponding shapes were defined
+        in the collision handler associated with this arbiter.
+
+        This is a shorthand to get the bodes::
+
+            arb.bodies == arb.shapes[0].body, arb.shapes[1].body .
+        """
+        a, b = self.shapes
+        assert (
+            a.body != None
+        ), "Shape should have a body. Could be a bug in Pymunk, please report"
+        assert (
+            b.body != None
+        ), "Shape should have a body. Could be a bug in Pymunk, please report"
+        return a.body, b.body
+
     @property
     def shapes(self) -> Tuple["Shape", "Shape"]:
         """Get the shapes in the order that they were defined in the

pymunk/tests/test_arbiter.py

@@ -300,7 +300,7 @@ def separate2(arb: p.Arbiter, space: p.Space, data: Any) -> None:
 
         self.assertTrue(self.called2)
 
-    def testShapes(self) -> None:
+    def testShapesAndBodies(self) -> None:
         s = p.Space()
         s.gravity = 0, -100
 
@@ -324,6 +324,8 @@ def pre_solve(arb: p.Arbiter, space: p.Space, data: Any) -> bool:
             self.assertEqual(len(arb.shapes), 2)
             self.assertEqual(arb.shapes[0], c1)
             self.assertEqual(arb.shapes[1], c2)
+            self.assertEqual(arb.bodies[0], arb.shapes[0].body)
+            self.assertEqual(arb.bodies[1], arb.shapes[1].body)
             return True
 
         s.add_collision_handler(1, 2).pre_solve = pre_solve