WIP Removed data from coll callbacks, made CollisionHandler private

Author: viblo

pymunk/_callbacks.py

@@ -203,23 +203,23 @@ def ext_cpCollisionBeginFunc(
     _arb: ffi.CData, _space: ffi.CData, data: ffi.CData
 ) -> None:
     handler = ffi.from_handle(data)
-    handler._begin(Arbiter(_arb, handler._space), handler._space, handler.data)
+    handler._begin(Arbiter(_arb, handler._space), handler._space)
 
 
 @ffi.def_extern()
 def ext_cpCollisionPreSolveFunc(
     _arb: ffi.CData, _space: ffi.CData, data: ffi.CData
 ) -> None:
     handler = ffi.from_handle(data)
-    handler._pre_solve(Arbiter(_arb, handler._space), handler._space, handler.data)
+    handler._pre_solve(Arbiter(_arb, handler._space), handler._space)
 
 
 @ffi.def_extern()
 def ext_cpCollisionPostSolveFunc(
     _arb: ffi.CData, _space: ffi.CData, data: ffi.CData
 ) -> None:
     handler = ffi.from_handle(data)
-    handler._post_solve(Arbiter(_arb, handler._space), handler._space, handler.data)
+    handler._post_solve(Arbiter(_arb, handler._space), handler._space)
 
 
 @ffi.def_extern()
@@ -234,7 +234,7 @@ def ext_cpCollisionSeparateFunc(
         # this try is needed since a separate callback will be called
         # if a colliding object is removed, regardless if its in a
         # step or not. Meaning the unlock must succeed
-        handler._separate(Arbiter(_arb, handler._space), handler._space, handler.data)
+        handler._separate(Arbiter(_arb, handler._space), handler._space)
     finally:
         handler._space._locked = orig_locked
 

pymunk/collision_handler.py

@@ -8,7 +8,7 @@
 from ._chipmunk_cffi import ffi, lib
 from .arbiter import Arbiter
 
-_CollisionCallback = Callable[[Arbiter, "Space", dict[Any, Any]], None]
+_CollisionCallback = Callable[[Arbiter, "Space"], None]
 
 
 class CollisionHandler(object):
@@ -149,7 +149,7 @@ def separate(self, func: _CollisionCallback) -> None:
             self._handler.separateFunc = lib.ext_cpCollisionSeparateFunc
 
     @staticmethod
-    def do_nothing(arbiter: Arbiter, space: "Space", data: dict[Any, Any]) -> None:
+    def do_nothing(arbiter: Arbiter, space: "Space") -> None:
         """The default do nothing method used for the post_solve and seprate
         callbacks.
 

pymunk/space.py

@@ -28,14 +28,15 @@
 _AddableObjects = Union[Body, Shape, Constraint]
 
 
-class Handlers(Mapping[Union[None, int, tuple[int, int]], CollisionHandler]):
+class HandlersMapping(Mapping[Union[None, int, tuple[int, int]], CollisionHandler]):
 
     def __init__(self, space: "Space") -> None:
         self.space = space
 
-    _handlers: dict[Union[None, int, tuple[int, int]], CollisionHandler] = {}
+        self._handlers: dict[Union[None, int, tuple[int, int]], CollisionHandler] = {}
 
     def __getitem__(self, key: Union[None, int, tuple[int, int]]) -> CollisionHandler:
+
         if key in self._handlers:
             return self._handlers[key]
         if key == None:
@@ -44,7 +45,7 @@ def __getitem__(self, key: Union[None, int, tuple[int, int]]) -> CollisionHandle
         elif isinstance(key, int):
             self._handlers[key] = self.space.add_collision_handler(key, None)
             return self._handlers[key]
-        elif isinstance(key, tuple):
+        elif isinstance(key, tuple) and len(key) == 2:
             assert isinstance(key, tuple)
             self._handlers[key] = self.space.add_collision_handler(key[0], key[1])
             return self._handlers[key]
@@ -166,7 +167,7 @@ def spacefree(cp_space: ffi.CData) -> None:
         self._remove_later: set[_AddableObjects] = set()
         self._bodies_to_check: set[Body] = set()
 
-        self._collision_handlers = Handlers(self)
+        self._collision_handlers = HandlersMapping(self)
 
     @property
     def shapes(self) -> KeysView[Shape]:
@@ -618,24 +619,16 @@ def step(self, dt: float) -> None:
 
         self._post_step_callbacks.clear()
 
-    @property
-    def collision_handlers(
-        self,
-    ) -> Mapping[Union[None, int, tuple[int, int]], CollisionHandler]:
-        return self.collision_handlers
-
-    def add_collision_handler(
+    def set_collision_callbacks(
         self,
         collision_type_a: Optional[int] = None,
         collision_type_b: Optional[int] = None,
         begin: Optional[_CollisionCallback] = None,
-        pre_step: Optional[_CollisionCallback] = None,
-        post_step: Optional[_CollisionCallback] = None,
+        pre_solve: Optional[_CollisionCallback] = None,
+        post_solve: Optional[_CollisionCallback] = None,
         separate: Optional[_CollisionCallback] = None,
-        data: Optional[Mapping] = None,
-    ) -> CollisionHandler:
-        """Return the :py:class:`CollisionHandler` for collisions between
-        objects of type collision_type_a and collision_type_b.
+    ):
+        """Set callbacks that will be called during the 4 phases of collision handling.
 
         Use None to indicate any collision_type.
 
@@ -650,6 +643,9 @@ def add_collision_handler(
         If multiple handlers match the collision, the order will be that the
         most specific handler is called first.
 
+        Note that if a handler already exist for the a,b pair, that existing
+        handler will be returned.
+
         :param int collision_type_a: Collision type a
         :param int collision_type_b: Collision type b
 
@@ -662,23 +658,32 @@ def add_collision_handler(
             collision_type_b, collision_type_a = collision_type_a, collision_type_b
 
         key = collision_type_a, collision_type_b
-        if key in self._handlers:
-            return self._handlers[key]
-
-        # CP_WILDCARD_COLLISION_TYPE
-        wildcard = int(ffi.cast("uintptr_t", ~0))
-        if collision_type_a == None:
-            collision_type_a = wildcard
+        if key not in self._handlers:
+            # CP_WILDCARD_COLLISION_TYPE
+            wildcard = int(ffi.cast("uintptr_t", ~0))
+            if collision_type_a == None:
+                collision_type_a = wildcard
 
-        if collision_type_b == None:
-            collision_type_b = wildcard
+            if collision_type_b == None:
+                collision_type_b = wildcard
 
-        h = lib.cpSpaceAddCollisionHandler(
-            self._space, collision_type_a, collision_type_b
-        )
-        ch = CollisionHandler(h, self)
-        self._handlers[key] = ch
-        return ch
+            h = lib.cpSpaceAddCollisionHandler(
+                self._space, collision_type_a, collision_type_b
+            )
+            ch = CollisionHandler(h, self)
+            self._handlers[key] = ch
+        else:
+            ch = self._handlers[key]
+
+        if begin != None:
+            ch.begin = begin
+        if pre_solve != None:
+            ch.pre_solve = pre_solve
+        if post_solve != None:
+            ch.post_solve = post_solve
+        if separate != None:
+            ch.separate = separate
+        return
 
     def add_post_step_callback(
         self,

pymunk/tests/test_arbiter.py

@@ -24,11 +24,11 @@ def testRestitution(self) -> None:
 
         s.add(b1, c1, b2, c2)
 
-        def pre_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve(arb: p.Arbiter, space: p.Space) -> None:
             self.assertEqual(arb.restitution, 0.18)
             arb.restitution = 1
 
-        s.add_collision_handler(1, 2).pre_solve = pre_solve
+        s.set_collision_callbacks(1, 2, pre_solve=pre_solve)
 
         for x in range(10):
             s.step(0.1)
@@ -52,11 +52,11 @@ def testFriction(self) -> None:
 
         s.add(b1, c1, b2, c2)
 
-        def pre_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve(arb: p.Arbiter, space: p.Space) -> None:
             self.assertEqual(arb.friction, 0.18)
             arb.friction = 1
 
-        s.add_collision_handler(1, 2).pre_solve = pre_solve
+        s.set_collision_callbacks(1, 2, pre_solve=pre_solve)
 
         for x in range(10):
             s.step(0.1)
@@ -80,14 +80,14 @@ def testSurfaceVelocity(self) -> None:
 
         s.add(b1, c1, b2, c2)
 
-        def pre_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve(arb: p.Arbiter, space: p.Space) -> None:
             self.assertAlmostEqual(arb.surface_velocity.x, 1.38461538462)
             self.assertAlmostEqual(arb.surface_velocity.y, -0.923076923077)
 
             arb.surface_velocity = (10, 10)
             # TODO: add assert check that setting surface_velocity has any effect
 
-        s.add_collision_handler(1, 2).pre_solve = pre_solve
+        s.set_collision_callbacks(1, 2, pre_solve=pre_solve)
         for x in range(5):
             s.step(0.1)
 
@@ -106,7 +106,7 @@ def testContactPointSet(self) -> None:
 
         s.add(b1, c1, b2, c2)
 
-        def pre_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve(arb: p.Arbiter, space: p.Space) -> None:
             # check inital values
             ps = arb.contact_point_set
             self.assertEqual(len(ps.points), 1)
@@ -142,7 +142,7 @@ def f() -> None:
 
             self.assertRaises(Exception, f)
 
-        s.add_collision_handler(2, 1).pre_solve = pre_solve
+        s.set_collision_callbacks(2, 1, pre_solve=pre_solve)
 
         s.step(0.1)
 
@@ -165,12 +165,12 @@ def testImpulse(self) -> None:
 
         self.post_solve_done = False
 
-        def post_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def post_solve(arb: p.Arbiter, space: p.Space) -> None:
             self.assertAlmostEqual(arb.total_impulse.x, 3.3936651583)
             self.assertAlmostEqual(arb.total_impulse.y, 4.3438914027)
             self.post_solve_done = True
 
-        s.add_collision_handler(1, 2).post_solve = post_solve
+        s.set_collision_callbacks(1, 2, post_solve=post_solve)
 
         s.step(0.1)
 
@@ -194,10 +194,10 @@ def testTotalKE(self) -> None:
         s.add(b1, c1, b2, c2)
         r = {}
 
-        def post_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def post_solve(arb: p.Arbiter, space: p.Space) -> None:
             r["ke"] = arb.total_ke
 
-        s.add_collision_handler(1, 2).post_solve = post_solve
+        s.set_collision_callbacks(1, 2, post_solve=post_solve)
 
         s.step(0.1)
 
@@ -220,17 +220,17 @@ def testIsFirstContact(self) -> None:
 
         s.add(b1, c1, b2, c2)
 
-        def pre_solve1(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve1(arb: p.Arbiter, space: p.Space) -> None:
             self.assertTrue(arb.is_first_contact)
 
-        s.add_collision_handler(1, 2).pre_solve = pre_solve1
+        s.set_collision_callbacks(1, 2, pre_solve=pre_solve1)
 
         s.step(0.1)
 
-        def pre_solve2(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve2(arb: p.Arbiter, space: p.Space) -> None:
             self.assertFalse(arb.is_first_contact)
 
-        s.add_collision_handler(1, 2).pre_solve = pre_solve2
+        s.set_collision_callbacks(1, 2, pre_solve=pre_solve2)
 
         s.step(0.1)
 
@@ -248,10 +248,10 @@ def testNormal(self) -> None:
         s.add(b1, c1, c2)
         r = {}
 
-        def pre_solve1(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve1(arb: p.Arbiter, space: p.Space) -> None:
             r["n"] = Vec2d(*arb.normal)
 
-        s.add_collision_handler(1, 2).pre_solve = pre_solve1
+        s.set_collision_callbacks(1, 2, pre_solve=pre_solve1)
 
         s.step(0.1)
 
@@ -277,11 +277,11 @@ def testIsRemoval(self) -> None:
 
         self.called1 = False
 
-        def separate1(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def separate1(arb: p.Arbiter, space: p.Space) -> None:
             self.called1 = True
             self.assertFalse(arb.is_removal)
 
-        s.add_collision_handler(1, 2).separate = separate1
+        s.set_collision_callbacks(1, 2, separate=separate1)
 
         for x in range(10):
             s.step(0.1)
@@ -292,11 +292,11 @@ def separate1(arb: p.Arbiter, space: p.Space, data: Any) -> None:
 
         self.called2 = False
 
-        def separate2(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def separate2(arb: p.Arbiter, space: p.Space) -> None:
             self.called2 = True
             self.assertTrue(arb.is_removal)
 
-        s.add_collision_handler(1, 2).separate = separate2
+        s.set_collision_callbacks(1, 2, separate=separate2)
         s.remove(b1, c1)
 
         self.assertTrue(self.called2)
@@ -320,15 +320,15 @@ def testShapesAndBodies(self) -> None:
 
         self.called = False
 
-        def pre_solve(arb: p.Arbiter, space: p.Space, data: Any) -> None:
+        def pre_solve(arb: p.Arbiter, space: p.Space) -> None:
             self.called = True
             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)
 
-        s.add_collision_handler(1, 2).post_solve = pre_solve
+        s.set_collision_callbacks(1, 2, post_solve=pre_solve)
 
         s.step(0.1)
         self.assertTrue(self.called)
@@ -408,15 +408,15 @@ def callback(
             # print("process_values, expected calls", process_values, expected_calls)
 
             s = setup()
-            h = s.add_collision_handler(1, 2)
+            h = s.set_collision_callbacks(1, 2)
             h.data["process_values"] = process_values
             h.data["expected"] = expected_calls
             h.data["result"] = []
 
-            h.begin = functools.partial(callback, "b")
-            h.pre_solve = functools.partial(callback, "p")
-            h.post_solve = functools.partial(callback, "t")
-            h.separate = functools.partial(callback, "s")
+            begin = functools.partial(callback, "b")
+            pre_solve = functools.partial(callback, "p")
+            post_solve = functools.partial(callback, "t")
+            separate = functools.partial(callback, "s")
 
             s.step(0.1)
             s.step(0.1)