Merge pull request #1364 from gabrielebndn/topic/pickleFrame

Pickle frame
stack-of-tasks · Jan 5, 2021 · 8f040db · 8f040db
2 parents 0ad63c1 + 5c98d84
commit 8f040db
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 2 deletions.
diff --git a/bindings/python/multibody/frame.hpp b/bindings/python/multibody/frame.hpp
@@ -23,6 +23,8 @@ namespace pinocchio
       void visit(PyClass& cl) const 
       {
         cl
+          .def(bp::init<>(bp::arg("self"),"Default constructor"))
+          .def(bp::init<const Frame &>(bp::args("self","other"),"Copy constructor"))
           .def(bp::init< const std::string&,const JointIndex, const FrameIndex, const SE3&,FrameType> ((bp::arg("name (string)"),bp::arg("index of parent joint"), bp::args("index of parent frame"), bp::arg("SE3 placement"), bp::arg("type (FrameType)")),
                 "Initialize from name, parent joint id, parent frame id and placement wrt parent joint."))
 
@@ -33,6 +35,8 @@ namespace pinocchio
                          &Frame::placement,
                          "placement in the parent joint local frame")
           .def_readwrite("type", &Frame::type, "type of the frame")
+          .def(bp::self == bp::self)
+          .def(bp::self != bp::self)
           ;
       }
 
@@ -54,9 +58,32 @@ namespace pinocchio
         .def(FramePythonVisitor())
         .def(CopyableVisitor<Frame>())
         .def(PrintableVisitor<Frame>())
+        .def_pickle(Pickle())
         ;
       }
 
+    private:
+      struct Pickle : bp::pickle_suite
+      {
+        static bp::tuple getinitargs(const Frame &)
+        {
+          return bp::make_tuple();
+        }
+
+        static bp::tuple getstate(const Frame & f)
+        {
+          return bp::make_tuple(f.name, f.parent, f.previousFrame, f.placement, (int)f.type);
+        }
+
+        static void setstate(Frame & f, bp::tuple tup)
+        {
+          f.name = bp::extract<std::string>(tup[0]); 
+          f.parent = bp::extract<JointIndex>(tup[1]); 
+          f.previousFrame = bp::extract<JointIndex>(tup[2]); 
+          f.placement = bp::extract<SE3&>(tup[3]); 
+          f.type = (FrameType)(int)bp::extract<int>(tup[4]); 
+        }
+      };
     };
 
 

diff --git a/src/multibody/frame.hpp b/src/multibody/frame.hpp
@@ -63,7 +63,9 @@ namespace pinocchio
     {}
 
     ///
-    /// \returns true if *this and other matches and have the same parent, name and type.
+    /// \brief Equality comparison operator.
+    ///
+    /// \returns true if *this is equal to other.
     ///
     /// \param[in] other The frame to which the current frame is compared.
     ///
@@ -76,7 +78,16 @@ namespace pinocchio
       && placement == other.placement
       && type == other.type ;
     }
-
+
+    ///
+    /// \returns true if *this is NOT equal to other.
+    ///
+    template<typename S2, int O2>
+    bool operator != (const FrameTpl<S2,O2> & other) const
+    {
+      return !(*this == other);
+    }
+
     /// \returns An expression of *this with the Scalar type casted to NewScalar.
     template<typename NewScalar>
     FrameTpl<NewScalar,Options> cast() const

diff --git a/unittest/python/bindings_frame.py b/unittest/python/bindings_frame.py
@@ -44,6 +44,30 @@ def test_placement_get_set(self):
         f.placement = new_placement
         self.assertTrue(np.allclose(f.placement.homogeneous, new_placement.homogeneous))
 
+    def test_frame_equality(self):
+        M = pin.SE3.Random()
+        frame1 = pin.Frame("name", 1, 2, M, pin.OP_FRAME)
+        frame2 = pin.Frame("name", 1, 2, M, pin.OP_FRAME)
+        frame3 = pin.Frame("othername", 3, 4, pin.SE3.Random(), pin.BODY)
+
+        self.assertTrue(frame1 == frame2)
+        self.assertFalse(frame1 != frame2)
+        self.assertTrue(frame1 != frame3)
+        self.assertFalse(frame1 == frame3)
+
+    def test_pickle(self):
+        import pickle
+
+        frame = pin.Frame("name", 1, 2, pin.SE3.Random(), pin.OP_FRAME)
+        filename = "frame.pickle"
+        with open(filename, 'wb') as f:
+            pickle.dump(frame, f)
+
+        with open(filename, 'rb') as f:
+            frame_copy = pickle.load(f)
+
+        self.assertEqual(frame, frame_copy)
+
     def test_getters(self):
         data = self.model.createData()
         q = pin.randomConfiguration(self.model)