Dump and load point clouds passing tests

idefix/io.py

@@ -169,9 +169,9 @@ def _arr_to_rec(arr):
 def load_pc(fname):
     """Load point cloud from file.
     
-    Loader for point clouds containted in '.pc', '.pcz' or compatible '.npy',
-    '.npz' files. This "point cloud" format is based on NumPy files, with small
-    overhead to manage record array and multispectral point clouds.
+    Loader for point clouds containted in compatible '.npz' files. This "point
+    cloud" format is based on NumPy files, with small overhead to manage record
+    array and multispectral point clouds.
 
     Parameters
     ----------
@@ -181,8 +181,8 @@ def load_pc(fname):
     Returns
     -------
     point_cloud : recarray or tuple of recarray
-        The point cloud respecting or tuple of point clouds (for multispectral
-        point cloud files).
+        The point cloud or tuple of point clouds (for multispectral point cloud
+        files).
     """
     log.info('Loading point cloud file \'{}\')'.format(fname))
 
@@ -193,3 +193,25 @@ def load_pc(fname):
         return _arr_to_rec(archive[archive.files[0]])
     else:
         return tuple(_arr_to_rec(archive[arr]) for arr in archive.files)
+
+def dump_pc(fname, point_cloud, compress=False):
+    """Dump point cloud to file.
+    
+    Write a point cloud (or several point clouds) in a '.npz' files.
+
+    Parameters
+    ----------
+    fname : str, Path
+        Path to the point cloud file to create.
+    point_cloud : recarray or tuple of recarray
+        The point cloud (or the tuple of point clouds) to dump.
+    compress : bool
+        Enable compression of the dumped file. Default is False.
+    """
+    if hasattr(point_cloud, 'spatial'):
+        point_cloud = (point_cloud, )
+
+    if compress:
+        np.savez_compressed(fname, *point_cloud)
+    else:
+        np.savez(fname, *point_cloud)

pytest.ini

@@ -0,0 +1 @@
+test/pytest.ini

test/pytest.ini

@@ -4,4 +4,5 @@ filterwarnings =
     ignore::DeprecationWarning:apptools.*:
     ignore::DeprecationWarning:pyface.*:
     ignore::DeprecationWarning:traits.*:
+    ignore::DeprecationWarning:traitsui.*:
     ignore:.*escape sequence.*:DeprecationWarning

test/test_io.py

@@ -89,13 +89,15 @@ def test_load_txt(datadir, fname, head, separator, exp_point_count, exp_field_co
 @pytest.mark.parametrize('fname, exp_point_count, exp_field_count', [
     ('test.npz', 58629, 2, ),
     ('test_compressed.npz', 58629, 2,),
+    ('test_multi.npz', (100, 200), 2,),
+    ('test_multi_compressed.npz', (100, 200), 2,),
 ])
 def test_load_pc(datadir, fname, exp_point_count, exp_field_count):
     fname = datadir.join(fname)
     
     # Raise "No such file"
     with pytest.raises(IOError) as e_info:
-        io.load_pc('not_as_file.las')
+        io.load_pc('not_as_file.npz')
 
     # Open file without exception
     try:
@@ -103,12 +105,51 @@ def test_load_pc(datadir, fname, exp_point_count, exp_field_count):
     except IOError:
         pytest.fail('Opening legit file without exception')
 
+    if isinstance(exp_point_count, tuple):
+        assert isinstance(result, tuple), "Multi point cloud file should return tuple of point cloud"
+        result = result[0]
+        exp_point_count = exp_point_count[0]
+
     assert result.size == exp_point_count, "Return correct point count"
 
     assert result['spatial'].shape[-1] == 3, "Return ndarray with spatial field"
 
+    assert result.spatial.shape[-1] == 3, "Returned array is not a recarray"
+
     assert (result['spatial'] == result.spatial).all(), "Quick access with records array"
 
     assert len(result['feature'].dtype) == exp_field_count, "Return ndarray with attribute fields"
 
     assert result.spatial.dtype == np.float, "Dtype of spatial is np.float"
+
+@pytest.mark.parametrize('fname, compress', [
+    ('test.npz', False,),
+    ('test.npz', True,),
+    ('test_multi.npz', False,),
+    ('test_multi.npz', True,),
+])
+def test_dump_pc(datadir, fname, compress):
+    in_fname = datadir.join(fname)
+    pc = io.load_pc(in_fname)
+
+    out_fname = datadir / 'PYTEST_test.npz'
+
+    try: 
+        io.dump_pc(out_fname, pc, compress)
+    except IOError:
+        pytest.fail('Dump file without exception')
+
+    assert out_fname.exists(), 'The dump file was not created'
+
+    in_out_pc = io.load_pc(out_fname)
+
+    assert len(in_out_pc) == len(pc), 'Missmatch of dumped point cloud'
+
+    if isinstance(pc, tuple):
+        assert in_out_pc[0].spatial.shape == pc[0].spatial.shape, 'Missmatch of dumped point cloud'
+        assert in_out_pc[0].spatial.dtype == pc[0].spatial.dtype, 'Missmatch of dumped point cloud'
+        assert in_out_pc[0].feature.dtype == pc[0].feature.dtype, 'Missmatch of dumped point cloud'
+    else:
+        assert in_out_pc.spatial.shape == pc.spatial.shape, 'Missmatch of dumped point cloud'
+        assert in_out_pc.spatial.dtype == pc.spatial.dtype, 'Missmatch of dumped point cloud'
+        assert in_out_pc.feature.dtype == pc.feature.dtype, 'Missmatch of dumped point cloud'