Reverse dependency between shape and body, so shape has the weak ref #275

Author: viblo

pymunk/body.py

@@ -224,7 +224,7 @@ def freebody(cp_body: ffi.CData) -> None:
         self._constraints: WeakSet["Constraint"] = (
             WeakSet()
         )  # weak refs to any constraints attached
-        self._shapes: WeakSet["Shape"] = WeakSet()  # weak refs to any shapes attached
+        self._shapes: dict["Shape", None] = {}
 
         d = ffi.new_handle(self)
         self._data_handle = d  # to prevent gc to collect the handle
@@ -660,10 +660,7 @@ def constraints(self) -> Set["Constraint"]:
 
     @property
     def shapes(self) -> Set["Shape"]:
-        """Get the shapes attached to this body.
-
-        The body only keeps a weak reference to the shapes and a live
-        body wont prevent GC of the attached shapes"""
+        """Get the shapes attached to this body."""
         return set(self._shapes)
 
     def local_to_world(self, v: Tuple[float, float]) -> Vec2d:

pymunk/shapes.py

@@ -1,5 +1,6 @@
 __docformat__ = "reStructuredText"
 
+import weakref
 from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple
 
 if TYPE_CHECKING:
@@ -42,18 +43,19 @@ class Shape(PickleMixin, TypingAttrMixing, object):
     _pickle_attrs_skip = PickleMixin._pickle_attrs_skip + ["mass", "density"]
 
     _space = None  # Weak ref to the space holding this body (if any)
-
-    _id_counter = 1
+    _dead_ref = weakref.ref(set())
 
     def __init__(self, shape: "Shape") -> None:
         self._shape = shape
-        self._body: Optional["Body"] = shape.body
+        self._body: weakref.ref = weakref.ref(shape.body)
 
     def _init(self, body: Optional["Body"], _shape: ffi.CData) -> None:
-        self._body = body
 
         if body is not None:
-            body._shapes.add(self)
+            self._body = weakref.ref(body)
+            body._shapes[self] = None
+        else:
+            self._body = Shape._dead_ref
 
         def shapefree(cp_shape: ffi.CData) -> None:
             cp_space = cp.cpShapeGetSpace(cp_shape)
@@ -225,19 +227,25 @@ def _set_surface_velocity(self, surface_v: Tuple[float, float]) -> None:
 
     @property
     def body(self) -> Optional["Body"]:
-        """The body this shape is attached to. Can be set to None to
-        indicate that this shape doesnt belong to a body."""
-        return self._body
+        """The body this shape is attached to.
+
+        Can be set to None to indicate that this shape doesnt belong to a body.
+        The shape only holds a weakref to the Body, meaning it wont prevent it
+        from being GCed.
+        """
+        return self._body()
 
     @body.setter
     def body(self, body: Optional["Body"]) -> None:
-        if self._body is not None:
-            self._body._shapes.remove(self)
-        body_body = ffi.NULL if body is None else body._body
-        cp.cpShapeSetBody(self._shape, body_body)
+        if self.body is not None:
+            del self.body._shapes[self]
+        cp_body = ffi.NULL if body is None else body._body
+        cp.cpShapeSetBody(self._shape, cp_body)
         if body is not None:
-            body._shapes.add(self)
-        self._body = body
+            body._shapes[self] = None
+            self._body = weakref.ref(body)
+        else:
+            self._body = Shape._dead_ref
 
     def update(self, transform: Transform) -> BB:
         """Update, cache and return the bounding box of a shape with an