From bcabd197962f4ed6e1d6941a403408c260365a2b Mon Sep 17 00:00:00 2001
From: Anton Soloviev <anton@praviel.com>
Date: Sun, 11 Oct 2026 12:12:41 -0700
Subject: [PATCH] Preserve custom spring force and torque callbacks when
 copying

---
 pymunk/constraints.py           |  4 +++
 pymunk/tests/test_constraint.py | 61 +++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/pymunk/constraints.py b/pymunk/constraints.py
index 0aafa1b1..ef090d7b 100644
--- a/pymunk/constraints.py
+++ b/pymunk/constraints.py
@@ -578,6 +578,8 @@ class DampedSpring(Constraint):
         "damping",
     ]
 
+    _pickle_attrs_general = Constraint._pickle_attrs_general + ["force_func"]
+
     _pickle_attrs_skip = Constraint._pickle_attrs_skip + ["_force_func"]
 
     def __init__(
@@ -707,6 +709,8 @@ class DampedRotarySpring(Constraint):
         "damping",
     ]
 
+    _pickle_attrs_general = Constraint._pickle_attrs_general + ["torque_func"]
+
     _pickle_attrs_skip = Constraint._pickle_attrs_skip + ["_torque_func"]
 
     _torque_func: Optional[_TorqueFunc] = None
diff --git a/pymunk/tests/test_constraint.py b/pymunk/tests/test_constraint.py
index 9c4873f4..21430e6a 100644
--- a/pymunk/tests/test_constraint.py
+++ b/pymunk/tests/test_constraint.py
@@ -1,3 +1,4 @@
+import copy
 import gc
 import pickle
 import unittest
@@ -418,6 +419,32 @@ def test_pickle(self) -> None:
         self.assertEqual(j.a.mass, j2.a.mass)
         self.assertEqual(j.b.mass, j2.b.mass)
 
+    def test_copy_force_func(self) -> None:
+        a, b = p.Body(10, 10), p.Body(20, 20)
+        b.position = (3, 0)
+        j = DampedSpring(a, b, (0, 0), (0, 0), 1, 2, 0)
+        j.force_func = nonlinear_spring_force
+        for clone in (
+            j.copy,
+            lambda: copy.deepcopy(j),
+            lambda: pickle.loads(pickle.dumps(j)),
+        ):
+            with self.subTest(clone=clone):
+                j2 = clone()
+                self.assertIs(j2.force_func, nonlinear_spring_force)
+                space = p.Space()
+                space.add(j2.a, j2.b, j2)
+                space.step(0.125)
+                self.assertEqual(j2.impulse, -2)
+
+        j.force_func = DampedSpring.spring_force
+        j2 = pickle.loads(pickle.dumps(j))
+        self.assertIs(j2.force_func, DampedSpring.spring_force)
+        space = p.Space()
+        space.add(j2.a, j2.b, j2)
+        space.step(0.125)
+        self.assertEqual(j2.impulse, -0.5)
+
 
 class UnitTestDampedRotarySpring(unittest.TestCase):
     def test_rest_angle(self) -> None:
@@ -478,6 +505,32 @@ def test_pickle(self) -> None:
         self.assertEqual(j.a.mass, j2.a.mass)
         self.assertEqual(j.b.mass, j2.b.mass)
 
+    def test_copy_torque_func(self) -> None:
+        a, b = p.Body(10, 10), p.Body(20, 20)
+        b.angle = 2
+        j = DampedRotarySpring(a, b, 0, 2, 0)
+        j.torque_func = nonlinear_spring_torque
+        for clone in (
+            j.copy,
+            lambda: copy.deepcopy(j),
+            lambda: pickle.loads(pickle.dumps(j)),
+        ):
+            with self.subTest(clone=clone):
+                j2 = clone()
+                self.assertIs(j2.torque_func, nonlinear_spring_torque)
+                space = p.Space()
+                space.add(j2.a, j2.b, j2)
+                space.step(0.125)
+                self.assertEqual(j2.impulse, -2)
+
+        j.torque_func = DampedRotarySpring.spring_torque
+        j2 = pickle.loads(pickle.dumps(j))
+        self.assertIs(j2.torque_func, DampedRotarySpring.spring_torque)
+        space = p.Space()
+        space.add(j2.a, j2.b, j2)
+        space.step(0.125)
+        self.assertEqual(j2.impulse, -0.5)
+
 
 class UnitTestRotaryLimitJoint(unittest.TestCase):
     def test_min(self) -> None:
@@ -597,3 +650,11 @@ def pre_solve(c: Constraint, s: p.Space) -> None:
 
 def post_solve(c: Constraint, s: p.Space) -> None:
     pass
+
+
+def nonlinear_spring_force(spring: DampedSpring, distance: float) -> float:
+    return spring.stiffness * (spring.rest_length - distance) ** 3
+
+
+def nonlinear_spring_torque(spring: DampedRotarySpring, relative_angle: float) -> float:
+    return spring.stiffness * (relative_angle - spring.rest_angle) ** 3