py mbp: Allow SpatialVector's to be pickled (RobotLocomotion#14104)

mcx · Sep 21, 2020 · 0223399 · 0223399
1 parent a216c5c
commit 0223399
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 2 deletions.
diff --git a/bindings/pydrake/multibody/BUILD.bazel b/bindings/pydrake/multibody/BUILD.bazel
@@ -252,7 +252,7 @@ drake_py_unittest(
     name = "math_test",
     deps = [
         ":math_py",
-        "//bindings/pydrake/common/test_utilities:numpy_compare_py",
+        "//bindings/pydrake/common/test_utilities",
     ],
 )
 

diff --git a/bindings/pydrake/multibody/math_py.cc b/bindings/pydrake/multibody/math_py.cc
@@ -77,7 +77,9 @@ void BindSpatialVectorMixin(PyClass* pcls) {
               This is done because defining ``__rmatmul__`` on this class does
               not disambiguate against the definitions of
               ``RotationMatrix.__matmul__``.
-          )""");
+          )""")
+      .def(py::pickle([](const Class& self) { return self.get_coeffs(); },
+          [](const Vector6<T>& coeffs) { return Class(coeffs); }));
   DefCopyAndDeepCopy(&cls);
 }
 

diff --git a/bindings/pydrake/multibody/test/math_test.py b/bindings/pydrake/multibody/test/math_test.py
@@ -6,6 +6,7 @@
 from pydrake.common.cpp_param import List
 from pydrake.common.value import Value
 import pydrake.common.test_utilities.numpy_compare as numpy_compare
+from pydrake.common.test_utilities.pickle_compare import assert_pickle
 from pydrake.symbolic import Expression
 from pydrake.math import RotationMatrix_
 from pydrake.multibody.math import (
@@ -75,6 +76,8 @@ def check_spatial_vector(
         R = RotationMatrix_[T]()
         numpy_compare.assert_float_equal((
             vec1.Rotate(R_FE=R)).get_coeffs(), coeffs_expected)
+        # Test pickling.
+        assert_pickle(self, vec1, cls.get_coeffs, T=T)
 
     @numpy_compare.check_all_types
     def test_spatial_vector_types(self, T):