Update body weakref to space to ref, and fix issue with del space. #278

Author: viblo

pymunk/_callbacks.py

@@ -290,8 +290,8 @@ def ext_cpBodyArbiterIteratorFunc(
     _body: ffi.CData, _arbiter: ffi.CData, data: ffi.CData
 ) -> None:
     body, func, args, kwargs = ffi.from_handle(data)
-    assert body._space is not None
-    arbiter = Arbiter(_arbiter, body._space)
+    assert body.space is not None
+    arbiter = Arbiter(_arbiter, body.space)
     func(arbiter, *args, **kwargs)
 
 

pymunk/body.py

@@ -1,5 +1,6 @@
 __docformat__ = "reStructuredText"
 
+import weakref
 from typing import (  # Literal,
     TYPE_CHECKING,
     Any,
@@ -106,12 +107,13 @@ class Body(PickleMixin, TypingAttrMixing, object):
         "is_sleeping",
         "_velocity_func",
         "_position_func",
+        # "_space",
     ]
 
     _position_func: Optional[_PositionFunc] = None
     _velocity_func: Optional[_VelocityFunc] = None
 
-    _id_counter = 1
+    _dead_ref = weakref.ref(set())
 
     def __init__(
         self, mass: float = 0, moment: float = 0, body_type: _BodyType = DYNAMIC
@@ -217,9 +219,7 @@ def freebody(cp_body: ffi.CData) -> None:
         elif body_type == Body.STATIC:
             self._body = ffi.gc(lib.cpBodyNewStatic(), freebody)
 
-        self._space: Optional["Space"] = (
-            None  # Weak ref to the space holding this body (if any)
-        )
+        self._space: weakref.ref = Body._dead_ref
 
         self._constraints: WeakSet["Constraint"] = (
             WeakSet()
@@ -264,7 +264,7 @@ def mass(self) -> float:
     @mass.setter
     def mass(self, mass: float) -> None:
         assert (
-            self._space is None or mass > 0
+            self.space is None or mass > 0
         ), "Dynamic bodies must have mass > 0 if they are attached to a Space."
         lib.cpBodySetMass(self._body, mass)
 
@@ -400,16 +400,14 @@ def rotation_vector(self) -> Vec2d:
     def space(self) -> Optional["Space"]:
         """Get the :py:class:`Space` that the body has been added to (or
         None)."""
-        assert hasattr(self, "_space"), (  # TODO: When can this happen?
+        # This assert is tested in test_pickle_circular_ref
+        assert hasattr(self, "_space"), (
             "_space not set. This can mean there's a direct or indirect"
             " circular reference between the Body and the Space. Circular"
             " references are not supported when using pickle or copy and"
             " might crash."
         )
-        if self._space is not None:
-            return self._space._get_self()  # ugly hack because of weakref
-        else:
-            return None
+        return self._space()
 
     def _set_velocity_func(self, func: _VelocityFunc) -> None:
         if func == Body.update_velocity:
@@ -571,7 +569,7 @@ def sleep(self) -> None:
 
         Cannot be called from a callback.
         """
-        if self._space == None:
+        if self.space == None:
             raise Exception("Body not added to space")
         lib.cpBodySleep(self._body)
 
@@ -589,7 +587,7 @@ def sleep_with_group(self, body: "Body") -> None:
         to initialize levels and start stacks of objects in a pre-sleeping
         state.
         """
-        if self._space == None:
+        if self.space == None:
             raise Exception("Body not added to space")
         lib.cpBodySleepWithGroup(self._body, body._body)
 

pymunk/space.py

@@ -167,7 +167,7 @@ def constraints(self) -> List[Constraint]:
         return list(self._constraints)
 
     def _setup_static_body(self, static_body: Body) -> None:
-        static_body._space = weakref.proxy(self)
+        static_body._space = weakref.ref(self)
         cp.cpSpaceAddBody(self._space, static_body._body)
 
     @property
@@ -407,7 +407,7 @@ def _add_body(self, body: "Body") -> None:
         assert body not in self._bodies, "Body already added to this space."
         assert body.space == None, "Body already added to another space."
 
-        body._space = weakref.proxy(self)
+        body._space = weakref.ref(self)
         self._bodies[body] = None
         self._bodies_to_check.add(body)
         cp.cpSpaceAddBody(self._space, body._body)
@@ -437,7 +437,7 @@ def _remove_shape(self, shape: "Shape") -> None:
     def _remove_body(self, body: "Body") -> None:
         """Removes a body from the space"""
         assert body in self._bodies, "body not in space, already removed?"
-        body._space = None
+        body._space = body._dead_ref
         if body in self._bodies_to_check:
             self._bodies_to_check.remove(body)
         # During GC at program exit sometimes the shape might already be removed. Then skip this step.

pymunk/tests/test_space.py

@@ -1168,6 +1168,24 @@ def testPickleCachedArbiters(self) -> None:
         self.assertAlmostEqual(s.bodies[0].position.x, s_copy.bodies[0].position.x)
         self.assertAlmostEqual(s.bodies[0].position.y, s_copy.bodies[0].position.y)
 
+    def testDeleteSpaceWithObjects(self) -> None:
+        s = p.Space()
+
+        b = p.Body(1)
+
+        c = p.Circle(b, 10)
+
+        j = p.PinJoint(b, s.static_body)
+
+        s.add(b, c, j)
+
+        del s
+
+        self.assertIsNone(b.space)
+        self.assertIsNone(c.space)
+        self.assertEqual(j.a, b)
+        self.assertEqual(j.b.body_type, p.Body.STATIC)
+
 
 def f1(*args: Any, **kwargs: Any) -> None:
     pass