Change Space.shapes internally to behave as Space.bodies

Author: viblo

pymunk/shapes.py

@@ -66,21 +66,8 @@ def shapefree(cp_shape: ffi.CData) -> None:
             cp.cpShapeFree(cp_shape)
 
         self._shape = ffi.gc(_shape, shapefree)
-        self._set_id()
-
-    @property
-    def _id(self) -> int:
-        """Unique id of the Shape.
-
-        .. note::
-            Experimental API. Likely to change in future major, minor orpoint
-            releases.
-        """
-        return int(ffi.cast("int", cp.cpShapeGetUserData(self._shape)))
-
-    def _set_id(self) -> None:
-        cp.cpShapeSetUserData(self._shape, ffi.cast("cpDataPointer", Shape._id_counter))
-        Shape._id_counter += 1
+        self._h = ffi.new_handle(self)  # to prevent GC of the handle
+        cp.cpShapeSetUserData(self._shape, self._h)
 
     @property
     def mass(self) -> float:
@@ -290,8 +277,8 @@ def point_query(self, p: Tuple[float, float]) -> PointQueryInfo:
         info = ffi.new("cpPointQueryInfo *")
         _ = cp.cpShapePointQuery(self._shape, p, info)
 
-        ud = int(ffi.cast("int", cp.cpShapeGetUserData(info.shape)))
-        assert ud == self._id
+        shape = ffi.from_handle(cp.cpShapeGetUserData(info.shape))
+        assert shape == self, "This is a bug in Pymunk. Please report it."
         return PointQueryInfo(
             self,
             Vec2d(info.point.x, info.point.y),
@@ -311,8 +298,8 @@ def segment_query(
         info = ffi.new("cpSegmentQueryInfo *")
         r = cp.cpShapeSegmentQuery(self._shape, start, end, radius, info)
         if r:
-            ud = int(ffi.cast("int", cp.cpShapeGetUserData(info.shape)))
-            assert ud == self._id
+            shape = ffi.from_handle(cp.cpShapeGetUserData(info.shape))
+            assert shape == self, "This is a bug in Pymunk. Please report it."
             return SegmentQueryInfo(
                 self,
                 Vec2d(info.point.x, info.point.y),

pymunk/space.py

@@ -132,9 +132,9 @@ def spacefree(cp_space: ffi.CData) -> None:
         )  # To prevent the gc to collect the callbacks.
 
         self._post_step_callbacks: Dict[Any, Callable[["Space"], None]] = {}
-        self._removed_shapes: Dict[int, Shape] = {}
+        self._removed_shapes: Dict[Shape, None] = {}
 
-        self._shapes: Dict[int, Shape] = {}
+        self._shapes: Dict[Shape, None] = {}
         self._bodies: Dict[Body, None] = {}
         self._static_body: Optional[Body] = None
         self._constraints: Dict[Constraint, None] = {}
@@ -154,7 +154,7 @@ def shapes(self) -> List[Shape]:
 
         (includes both static and non-static)
         """
-        return list(self._shapes.values())
+        return list(self._shapes)
 
     @property
     def bodies(self) -> List[Body]:
@@ -389,8 +389,7 @@ def remove(self, *objs: _AddableObjects) -> None:
 
     def _add_shape(self, shape: "Shape") -> None:
         """Adds a shape to the space"""
-        # print("addshape", self._space, shape)
-        assert shape._id not in self._shapes, "Shape already added to space."
+        assert shape not in self._shapes, "Shape already added to space."
         assert (
             shape.space == None
         ), "Shape already added to another space. A shape can only be in one space at a time."
@@ -400,7 +399,7 @@ def _add_shape(self, shape: "Shape") -> None:
         ), "The shape's body must be added to the space before (or at the same time) as the shape."
 
         shape._space = weakref.proxy(self)
-        self._shapes[shape._id] = shape
+        self._shapes[shape] = None
         cp.cpSpaceAddShape(self._space, shape._shape)
 
     def _add_body(self, body: "Body") -> None:
@@ -427,13 +426,13 @@ def _add_constraint(self, constraint: "Constraint") -> None:
 
     def _remove_shape(self, shape: "Shape") -> None:
         """Removes a shape from the space"""
-        assert shape._id in self._shapes, "shape not in space, already removed?"
-        self._removed_shapes[shape._id] = shape
+        assert shape in self._shapes, "shape not in space, already removed?"
+        self._removed_shapes[shape] = None
         shape._space = None
         # During GC at program exit sometimes the shape might already be removed. Then skip this step.
         if cp.cpSpaceContainsShape(self._space, shape._shape):
             cp.cpSpaceRemoveShape(self._space, shape._shape)
-        del self._shapes[shape._id]
+        del self._shapes[shape]
 
     def _remove_body(self, body: "Body") -> None:
         """Removes a body from the space"""
@@ -517,7 +516,7 @@ def use_spatial_hash(self, dim: float, count: int) -> None:
         the shape to be inserted into many cells, setting it too low will
         cause too many objects into the same hash slot.
 
-        count is the suggested minimum number of cells in the hash table. If
+        count is the minimum number of cells in the hash table. If
         there are too few cells, the spatial hash will return many false
         positives. Too many cells will be hard on the cache and waste memory.
         Setting count to ~10x the number of objects in the space is probably a
@@ -563,7 +562,7 @@ def step(self, dt: float) -> None:
                 cp.cpHastySpaceStep(self._space, dt)
             else:
                 cp.cpSpaceStep(self._space, dt)
-            self._removed_shapes = {}
+            self._removed_shapes.clear()
         finally:
             self._locked = False
         self.add(*self._add_later)
@@ -575,7 +574,7 @@ def step(self, dt: float) -> None:
         for key in self._post_step_callbacks:
             self._post_step_callbacks[key](self)
 
-        self._post_step_callbacks = {}
+        self._post_step_callbacks.clear()
 
     def add_collision_handler(
         self, collision_type_a: int, collision_type_b: int
@@ -745,16 +744,7 @@ def point_query(
     def _get_shape(self, _shape: Any) -> Optional[Shape]:
         if not bool(_shape):
             return None
-
-        shapeid = int(ffi.cast("int", cp.cpShapeGetUserData(_shape)))
-        # return self._shapes[hashid_private]
-
-        if shapeid in self._shapes:
-            return self._shapes[shapeid]
-        elif shapeid in self._removed_shapes:
-            return self._removed_shapes[shapeid]
-        else:
-            return None
+        return ffi.from_handle(cp.cpShapeGetUserData(_shape))
 
     def point_query_nearest(
         self, point: Tuple[float, float], max_distance: float, shape_filter: ShapeFilter

pymunk/tests/test_shape.py

@@ -6,10 +6,6 @@
 
 
 class UnitTestShape(unittest.TestCase):
-    def testId(self) -> None:
-        c = p.Circle(None, 4)
-        self.assertGreater(c._id, 0)
-
     def testPointQuery(self) -> None:
         b = p.Body(10, 10)
         c = p.Circle(b, 5)