pcl.save for dumping to file; PCD and PLY format

3 files changed, 52 insertions, 27 deletions

diff --git a/pcl/__init__.py b/pcl/__init__.py
index ea680b2..f68e2f5 100644
--- a/pcl/__init__.py
+++ b/pcl/__init__.py
@@ -9,22 +9,39 @@ def load(path, format=None):
 
     Format should be "pcd", "ply", or None to infer from the pathname.
     """
-
-    if format is None:
-        for candidate in ["pcd", "ply"]:
-            if path.endswith("." + candidate):
-                format = candidate
-                break
-        else:
-            raise ValueError("Could not determine file format from pathname %s"
-                             % path)
-
+    format = _infer_format(path, format)
     p = PointCloud()
     try:
-        loader = getattr(p, "_from_%s_file" % format.lower())
+        loader = getattr(p, "_from_%s_file" % format)
     except AttributeError:
         raise ValueError("unknown file format %s" % format)
     if loader(path):
         raise IOError("error while loading pointcloud from %r (format=%r)"
                       % (path, format))
     return p
+
+
+def save(cloud, path, format=None, binary=False):
+    """Save pointcloud to file.
+
+    Format should be "pcd", "ply", or None to infer from the pathname.
+    """
+    format = _infer_format(path, format)
+    try:
+        dumper = getattr(cloud, "_to_%s_file" % format)
+    except AttributeError:
+        raise ValueError("unknown file format %s" % format)
+    if dumper(path, binary):
+        raise IOError("error while saving pointcloud to %r (format=%r)"
+                      % (path, format))
+
+
+def _infer_format(path, format):
+    if format is not None:
+        return format.lower()
+
+    for candidate in ["pcd", "ply"]:
+        if path.endswith("." + candidate):
+            return candidate
+
+    raise ValueError("Could not determine file format from pathname %s" % path)
diff --git a/pcl/_pcl.pyx b/pcl/_pcl.pyx
index e048906..c20df9f 100644
--- a/pcl/_pcl.pyx
+++ b/pcl/_pcl.pyx
@@ -240,31 +240,37 @@ cdef class PointCloud:
         return self._from_pcd_file(f)
 
     def _from_pcd_file(self, const char *s):
-        cdef int ok = 0
+        cdef int error = 0
         with nogil:
             ok = cpp.loadPCDFile(string(s), deref(self.thisptr))
-        return ok
+        return error
 
     def _from_ply_file(self, const char *s):
         cdef int ok = 0
         with nogil:
-            ok = cpp.loadPLYFile(string(s), deref(self.thisptr))
-        return ok
+            error = cpp.loadPLYFile(string(s), deref(self.thisptr))
+        return error
 
-    def to_file(self, char *f, bool ascii=True):
-        """
-        Save this pointcloud to a local file.
-        Only saving to binary or ascii pcd is supported
+    def to_file(self, const char *fname, bool ascii=True):
+        """Save pointcloud to a file in PCD format.
+
+        Deprecated: use pcl.save instead.
         """
-        cdef bool binary = not ascii
-        cdef int ok = 0
+        return self._to_pcd_file(fname, not ascii)
+
+    def _to_pcd_file(self, const char *f, bool binary=False):
+        cdef int error = 0
         cdef string s = string(f)
-        if f.endswith(".pcd"):
-            with nogil:
-                ok = cpp.savePCDFile(s, deref(self.thisptr), binary)
-        else:
-            raise ValueError("Incorrect file extension (must be .pcd)")
-        return ok
+        with nogil:
+            error = cpp.savePCDFile(s, deref(self.thisptr), binary)
+        return error
+
+    def _to_ply_file(self, const char *f, bool binary=False):
+        cdef int error = 0
+        cdef string s = string(f)
+        with nogil:
+            error = cpp.savePLYFile(s, deref(self.thisptr), binary)
+        return error
 
     def make_segmenter(self):
         """
diff --git a/pcl/pcl_defs.pxd b/pcl/pcl_defs.pxd
index ee2ce04..2a6abe8 100644
--- a/pcl/pcl_defs.pxd
+++ b/pcl/pcl_defs.pxd
@@ -128,6 +128,8 @@ cdef extern from "pcl/io/pcd_io.h" namespace "pcl::io":
 
 cdef extern from "pcl/io/ply_io.h" namespace "pcl::io":
     int loadPLYFile(string file_name, PointCloud[PointXYZ] cloud) nogil
+    int savePLYFile(string file_name, PointCloud[PointXYZ] cloud,
+                    bool binary_mode) nogil
 
 #http://dev.pointclouds.org/issues/624
 #cdef extern from "pcl/io/ply_io.h" namespace "pcl::io":