Preserve segment neighbor tangents when copying

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):

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)