diff --git a/pymunk/shapes.py b/pymunk/shapes.py index dcec992b..0173114c 100644 --- a/pymunk/shapes.py +++ b/pymunk/shapes.py @@ -435,6 +435,7 @@ class Segment(Shape): """ _pickle_attrs_init = Shape._pickle_attrs_init + ["a", "b", "radius"] + _neighbor_tangents: Optional[tuple[Vec2d, Vec2d]] = None def __init__( self, @@ -521,6 +522,19 @@ def set_neighbors( assert len(prev) == 2 assert len(next) == 2 cp.cpSegmentShapeSetNeighbors(self._shape, prev, next) + self._neighbor_tangents = (Vec2d(*prev) - self.a, Vec2d(*next) - self.b) + + def __getstate__(self) -> _State: + state = super().__getstate__() + if self._neighbor_tangents is not None: + state["special"].append(("_neighbor_tangents", self._neighbor_tangents)) + return state + + def __setstate__(self, state: _State) -> None: + super().__setstate__(state) + for key, value in state["special"]: + if key == "_neighbor_tangents": + self.set_neighbors(self.a + value[0], self.b + value[1]) class Poly(Shape): diff --git a/pymunk/tests/test_shape.py b/pymunk/tests/test_shape.py index 20bb4283..90369b17 100644 --- a/pymunk/tests/test_shape.py +++ b/pymunk/tests/test_shape.py @@ -256,6 +256,40 @@ def testPickle(self) -> None: class UnitTestSegment(unittest.TestCase): + def testCopyNeighbors(self) -> None: + for change_endpoints in (False, True): + space = p.Space() + segment = p.Segment(space.static_body, (-10, 0), (0, 0), 0) + segment.set_neighbors((-20, 0), (10, 0)) + if change_endpoints: + segment.unsafe_set_endpoints((-9, 0), (1, 0)) + space.add(segment) + circle = p.Circle(p.Body(1, 1), 0.5) + circle.body.position = segment.b + (0.25, 0.3) + self.assertEqual(len(space.shape_query(circle)), 0) + for restore in ( + lambda value: value.copy(), + lambda value: pickle.loads(pickle.dumps(value, 2)), + lambda value: pickle.loads(pickle.dumps(value, 5)), + ): + with self.subTest(change_endpoints=change_endpoints, restore=restore): + restored = restore(space) + self.assertEqual(len(restored.shape_query(circle)), 0) + copied_segment = restore(segment) + rebuilt = p.Space() + assert copied_segment.body is not None + rebuilt.add(copied_segment.body, copied_segment) + self.assertEqual(len(rebuilt.shape_query(circle)), 0) + + def testLegacySegmentState(self) -> None: + segment = p.Segment(None, (-10, 0), (0, 0), 0) + state = segment.__getstate__() + state["special"] = [] + restored = p.Segment.__new__(p.Segment) + restored.__setstate__(state) + self.assertEqual(restored.a, segment.a) + self.assertEqual(restored.b, segment.b) + def testBB(self) -> None: s = p.Space() b = p.Body(10, 10)