added copy interface

nmaxwell · Apr 25, 2018 · 73576ff · 73576ff
1 parent c2d2534
commit 73576ff
Show file tree

Hide file tree

Showing 3 changed files with 88 additions and 0 deletions.
diff --git a/src/Mesh.hh b/src/Mesh.hh
@@ -640,6 +640,13 @@ void expose_mesh(py::module& m, const char *_name) {
 				return mesh;
 			}), py::arg("points"), py::arg("face_vertex_indices")=py::array_t<int>())
 
+		//======================================================================
+		//  Copy interface
+		//======================================================================
+
+		.def("__copy__", &Mesh::py_copy)
+		.def("__deepcopy__", &Mesh::py_deepcopy)
+
 		//======================================================================
 		//  KernelT
 		//======================================================================

diff --git a/src/MeshTypes.hh b/src/MeshTypes.hh
@@ -174,6 +174,30 @@ public:
 		Mesh::copy_property(prop, _from, _to);
 	}
 
+	py::object py_copy() {
+		return py::cast(MeshWrapperT(*this));
+	}
+
+	py::object py_deepcopy(py::dict _memo) {
+		#if PY_MAJOR_VERSION < 3
+		py::object id = py::module::import("__builtin__").attr("id");
+		#else
+		py::object id = py::module::import("builtins").attr("id");
+		#endif
+		py::object deepcopy = py::module::import("copy").attr("deepcopy");
+
+		MeshWrapperT *copy = new MeshWrapperT(*this);
+		py::object copy_pyobj = py::cast(copy, py::return_value_policy::take_ownership);
+		_memo[id(py::cast(this))] = copy_pyobj;
+
+		py_deepcopy_prop<OpenMesh::VertexHandle>(copy, deepcopy, _memo);
+		py_deepcopy_prop<OpenMesh::HalfedgeHandle>(copy, deepcopy, _memo);
+		py_deepcopy_prop<OpenMesh::EdgeHandle>(copy, deepcopy, _memo);
+		py_deepcopy_prop<OpenMesh::FaceHandle>(copy, deepcopy, _memo);
+
+		return copy_pyobj;
+	}
+
 	size_t py_n_items(OpenMesh::VertexHandle) const { return Mesh::n_vertices(); }
 	size_t py_n_items(OpenMesh::HalfedgeHandle) const { return Mesh::n_halfedges(); }
 	size_t py_n_items(OpenMesh::EdgeHandle) const { return Mesh::n_edges(); }
@@ -186,6 +210,17 @@ public:
 
 private:
 
+	template <class Handle>
+	void py_deepcopy_prop(MeshWrapperT *_copy, py::object _copyfunc, py::dict _memo) {
+		for (const auto& item : py_prop_map(Handle())) {
+			const auto prop = item.second;
+			for (size_t i = 0; i < py_n_items(Handle()); ++i) {
+				const Handle h(i);
+				_copy->property(prop, h) = _copyfunc(this->property(prop, h), _memo);
+			}
+		}
+	}
+
 	template <class Handle, class PropHandle>
 	PropHandle py_prop_on_demand(const std::string& _name) {
 		auto& prop_map = py_prop_map(Handle());

diff --git a/tests/test_copy.py b/tests/test_copy.py
@@ -0,0 +1,46 @@
+import unittest
+import openmesh
+import copy
+
+class Python(unittest.TestCase):
+
+    def setUp(self):
+        self.mesh = openmesh.TriMesh()
+        self.vh1 = self.mesh.add_vertex([0, 0, 0])
+        self.vh2 = self.mesh.add_vertex([1, 0, 0])
+        self.vh3 = self.mesh.add_vertex([1, 1, 0])
+        self.fh = self.mesh.add_face(self.vh1, self.vh2, self.vh3)
+        self.mesh.set_vertex_property('test', self.vh1, self.mesh)
+        self.mesh.set_vertex_property('test', self.vh2, self.mesh)
+        self.one_two_three = [1, 2, 3]
+        self.mesh.set_vertex_property('test', self.vh3, self.one_two_three)
+        self.mesh.set_face_property('test', self.fh, self.one_two_three)
+
+    def test_shallowcopy(self):
+        mesh_copy = copy.copy(self.mesh)
+        mesh_copy.vertex_property('test', self.vh3)[:] = [4, 5, 6]
+        self.assertIsNot(self.mesh, mesh_copy)
+        self.assertIs(mesh_copy.vertex_property('test', self.vh1), self.mesh)
+        self.assertIs(mesh_copy.vertex_property('test', self.vh2), self.mesh)
+        self.assertIs(mesh_copy.vertex_property('test', self.vh3), self.one_two_three)
+        self.assertIs(mesh_copy.face_property('test', self.fh), self.one_two_three)
+        self.assertIs(mesh_copy.vertex_property('test', self.vh3), mesh_copy.face_property('test', self.fh))
+        self.assertEqual(self.mesh.vertex_property('test', self.vh3), [4, 5, 6])
+        self.assertEqual(mesh_copy.vertex_property('test', self.vh3), [4, 5, 6])
+
+    def test_deepcopy(self):
+        mesh_copy = copy.deepcopy(self.mesh)
+        mesh_copy.vertex_property('test', self.vh3)[:] = [4, 5, 6]
+        self.assertIsNot(self.mesh, mesh_copy)
+        self.assertIs(mesh_copy.vertex_property('test', self.vh1), mesh_copy)
+        self.assertIs(mesh_copy.vertex_property('test', self.vh2), mesh_copy)
+        self.assertIsNot(mesh_copy.vertex_property('test', self.vh3), self.one_two_three)
+        self.assertIsNot(mesh_copy.face_property('test', self.fh), self.one_two_three)
+        self.assertIs(mesh_copy.vertex_property('test', self.vh3), mesh_copy.face_property('test', self.fh))
+        self.assertListEqual(self.mesh.vertex_property('test', self.vh3), [1, 2, 3])
+        self.assertListEqual(mesh_copy.vertex_property('test', self.vh3), [4, 5, 6])
+
+
+if __name__ == '__main__':
+    suite = unittest.TestLoader().loadTestsFromTestCase(Python)
+    unittest.TextTestRunner(verbosity=2).run(suite)