implement copy functions

Author: alimanfoo

docs/api/convenience.rst

@@ -6,4 +6,7 @@ Convenience functions (``zarr.convenience``)
 .. autofunction:: load
 .. autofunction:: save_array
 .. autofunction:: save_group
+.. autofunction:: copy
+.. autofunction:: copy_all
 .. autofunction:: copy_store
+.. autofunction:: tree

zarr/__init__.py

@@ -11,5 +11,6 @@
 from zarr.hierarchy import group, open_group, Group
 from zarr.sync import ThreadSynchronizer, ProcessSynchronizer
 from zarr.codecs import *
-from zarr.convenience import open, save, save_array, save_group, load, copy_store
+from zarr.convenience import (open, save, save_array, save_group, load, copy_store,
+                              copy, copy_all, tree)
 from zarr.version import version as __version__

zarr/convenience.py

@@ -11,7 +11,7 @@
 from zarr.hierarchy import open_group, group as _create_group, Group
 from zarr.storage import contains_array, contains_group
 from zarr.errors import err_path_not_found
-from zarr.util import normalize_storage_path
+from zarr.util import normalize_storage_path, TreeViewer
 
 
 # noinspection PyShadowingBuiltins
@@ -507,16 +507,59 @@ def copy_store(source, dest, source_path='', dest_path='', excludes=None,
                 dest[dest_key] = source[source_key]
 
 
-def copy(source, dest, name=None, without_attrs=False, log=None, **create_kws):
-    """TODO"""
+def copy(source, dest, name=None, shallow=False, without_attrs=False, log=None,
+         **create_kws):
+    """Copy the source object into the given destination.
+
+    Parameters
+    ----------
+    source : group or array/dataset
+        A zarr group or array, or an h5py group or dataset.
+    dest : group
+        A zarr or h5py group.
+    name : str, optional
+        Name to copy the object to.
+    shallow : bool, optional
+        If True, only copy immediate children of `source`.
+    without_attrs : bool, optional
+        Do not copy user attributes.
+    log : callable, file path or file-like object, optional
+        If provided, will be used to log progress information.
+    **create_kws
+        Passed through to the create_dataset method when copying an array/dataset.
+
+    Examples
+    --------
+    >>> import h5py
+    >>> import zarr
+    >>> import numpy as np
+    >>> source = h5py.File('data/example.h5', mode='w')
+    >>> foo = source.create_group('foo')
+    >>> baz = foo.create_dataset('bar/baz', data=np.arange(100), chunks=(50,))
+    >>> spam = source.create_dataset('spam', data=np.arange(100, 200), chunks=(30,))
+    >>> zarr.tree(source)
+    /
+     ├── foo
+     │   └── bar
+     │       └── baz (100,) int64
+     └── spam (100,) int64
+    >>> dest = zarr.group()
+    >>> zarr.copy(source['foo'], dest)
+    >>> dest.tree()  # N.B., no spam
+    /
+     └── foo
+         └── bar
+             └── baz (100,) int64
+
+    """
 
     # setup logging
     with _LogWriter(log) as log:
-        _copy(log, source, dest, name=name, without_attrs=without_attrs, **create_kws)
+        _copy(log, source, dest, name=name, root=True, shallow=shallow,
+              without_attrs=without_attrs, **create_kws)
 
 
-def _copy(log, source, dest, name=None, without_attrs=False, **create_kws):
-    """TODO"""
+def _copy(log, source, dest, name, root, shallow, without_attrs, **create_kws):
 
     # are we copying to/from h5py?
     source_h5py = source.__module__.startswith('h5py.')
@@ -569,4 +612,119 @@ def _copy(log, source, dest, name=None, without_attrs=False, **create_kws):
         if not without_attrs:
             ds.attrs.update(source.attrs)
 
-    else:
+    elif root or not shallow:
+        # copy a group
+
+        # creat new group in destination
+        grp = dest.create_group(name)
+
+        # copy attributes
+        if not without_attrs:
+            grp.attrs.update(source.attrs)
+
+        # recurse
+        for k in source.keys():
+            _copy(log, source[k], grp, name=k, root=False, shallow=shallow,
+                  without_attrs=without_attrs, **create_kws)
+
+
+def tree(grp, expand=False, level=None):
+    """Provide a ``print``-able display of the hierarchy. This function is provided
+    mainly as a convenience for obtaining a tree view of an h5py group - zarr groups
+    have a ``.tree()`` method.
+
+    Parameters
+    ----------
+    grp : Group
+        Zarr or h5py group.
+    expand : bool, optional
+        Only relevant for HTML representation. If True, tree will be fully expanded.
+    level : int, optional
+        Maximum depth to descend into hierarchy.
+
+    Examples
+    --------
+    >>> import zarr
+    >>> g1 = zarr.group()
+    >>> g2 = g1.create_group('foo')
+    >>> g3 = g1.create_group('bar')
+    >>> g4 = g3.create_group('baz')
+    >>> g5 = g3.create_group('qux')
+    >>> d1 = g5.create_dataset('baz', shape=100, chunks=10)
+    >>> g1.tree()
+    /
+     ├── bar
+     │   ├── baz
+     │   └── qux
+     │       └── baz (100,) float64
+     └── foo
+    >>> import h5py
+    >>> h5f = h5py.File('data/example.h5', mode='w')
+    >>> zarr.copy_all(g1, h5f)
+    >>> zarr.tree(h5f)
+    /
+     ├── bar
+     │   ├── baz
+     │   └── qux
+     │       └── baz (100,) float64
+     └── foo
+
+    See Also
+    --------
+    zarr.hierarchy.Group.tree
+
+    """
+
+    return TreeViewer(grp, expand=expand, level=level)
+
+
+def copy_all(source, dest, shallow=False, without_attrs=False, log=None, **create_kws):
+    """Copy all children of the source group into the destination group.
+
+    Parameters
+    ----------
+    source : group or array/dataset
+        A zarr group or array, or an h5py group or dataset.
+    dest : group
+        A zarr or h5py group.
+    shallow : bool, optional
+        If True, only copy immediate children of `source`.
+    without_attrs : bool, optional
+        Do not copy user attributes.
+    log : callable, file path or file-like object, optional
+        If provided, will be used to log progress information.
+    **create_kws
+        Passed through to the create_dataset method when copying an array/dataset.
+
+    Examples
+    --------
+    >>> import h5py
+    >>> import zarr
+    >>> import numpy as np
+    >>> source = h5py.File('data/example.h5', mode='w')
+    >>> foo = source.create_group('foo')
+    >>> baz = foo.create_dataset('bar/baz', data=np.arange(100), chunks=(50,))
+    >>> spam = source.create_dataset('spam', data=np.arange(100, 200), chunks=(30,))
+    >>> zarr.tree(source)
+    /
+     ├── foo
+     │   └── bar
+     │       └── baz (100,) int64
+     └── spam (100,) int64
+    >>> dest = zarr.group()
+    >>> zarr.copy_all(source, dest)
+    >>> dest.tree()
+    /
+     ├── foo
+     │   └── bar
+     │       └── baz (100,) int64
+     └── spam (100,) int64
+
+    """
+
+    # setup logging
+    with _LogWriter(log) as log:
+        for k in source.keys():
+            _copy(log, source[k], dest, name=k, root=False, shallow=shallow,
+                  without_attrs=without_attrs, **create_kws)
+

zarr/tests/test_convenience.py

@@ -2,12 +2,14 @@
 from __future__ import absolute_import, print_function, division
 import tempfile
 import atexit
+import os
 
 
 from nose.tools import assert_raises
 import numpy as np
 from numpy.testing import assert_array_equal
 from numcodecs import Zlib
+import pytest
 
 
 from zarr.convenience import open, save, save_group, load, copy_store, copy
@@ -180,78 +182,135 @@ def test_copy_store():
     assert 'bar/qux' in dest
 
 
-def test_copy():
+def _test_copy(new_source, new_dest):
 
-    source = group()
+    source = new_source()
     foo = source.create_group('foo')
     foo.attrs['experiment'] = 'weird science'
-    baz = foo.create_dataset('bar/baz', data=np.arange(100), chunks=50)
+    baz = foo.create_dataset('bar/baz', data=np.arange(100), chunks=(50,))
     baz.attrs['units'] = 'metres'
+    spam = source.create_dataset('spam', data=np.arange(100, 200), chunks=(30,))
 
     # copy array with default options
-    dest = group()
+    dest = new_dest()
     copy(source['foo/bar/baz'], dest)
     a = dest['baz']  # defaults to use source name
-    assert isinstance(a, Array)
     assert a.dtype == baz.dtype
     assert a.shape == baz.shape
     assert a.chunks == baz.chunks
-    assert a.compressor == baz.compressor
+    if hasattr(a, 'compressor') and hasattr(baz, 'compressor'):
+        assert a.compressor == baz.compressor
     assert_array_equal(a[:], baz[:])
     assert a.attrs['units'] == 'metres'
 
     # copy array with name
-    dest = group()
+    dest = new_dest()
     copy(source['foo/bar/baz'], dest, name='qux')
     assert 'baz' not in dest
     a = dest['qux']
-    assert isinstance(a, Array)
     assert a.dtype == baz.dtype
     assert a.shape == baz.shape
     assert a.chunks == baz.chunks
-    assert a.compressor == baz.compressor
+    if hasattr(a, 'compressor') and hasattr(baz, 'compressor'):
+        assert a.compressor == baz.compressor
     assert_array_equal(a[:], baz[:])
     assert a.attrs['units'] == 'metres'
 
     # copy array, provide creation options
-    compressor = Zlib(1)
-    chunks = True
-    copy(source['foo/bar/baz'], dest, without_attrs=True, compressor=compressor,
-         chunks=chunks)
+    dest = new_dest()
+    compressor = Zlib(9)
+    if isinstance(dest, Group):
+        copy(source['foo/bar/baz'], dest, without_attrs=True, compressor=compressor,
+             chunks=True)
+    else:
+        copy(source['foo/bar/baz'], dest, without_attrs=True, compression='gzip',
+             compression_opts=9, chunks=True)
     a = dest['baz']
-    assert isinstance(a, Array)
     assert a.dtype == baz.dtype
     assert a.shape == baz.shape
     assert a.chunks != baz.chunks  # autochunking was requested
-    assert a.compressor == compressor
+    if hasattr(a, 'compressor'):
+        assert compressor == a.compressor
+        if hasattr(baz, 'compressor'):
+            assert a.compressor != baz.compressor
+    else:
+        assert a.compression == 'gzip'
+        assert a.compression_opts == 9
     assert_array_equal(a[:], baz[:])
     assert 'units' not in a.attrs
 
     # copy group, default options
-    dest = group()
+    dest = new_dest()
     copy(source['foo'], dest)
     g = dest['foo']  # defaults to use source name
-    assert isinstance(g, Group)
     assert g.attrs['experiment'] == 'weird science'
     a = g['bar/baz']
     assert a.dtype == baz.dtype
     assert a.shape == baz.shape
     assert a.chunks == baz.chunks
-    assert a.compressor == baz.compressor
+    if hasattr(a, 'compressor') and hasattr(baz, 'compressor'):
+        assert a.compressor == baz.compressor
     assert_array_equal(a[:], baz[:])
     assert a.attrs['units'] == 'metres'
 
     # copy group, non-default options
-    dest = group()
+    dest = new_dest()
     copy(source['foo'], dest, name='qux', without_attrs=True)
     assert 'foo' not in dest
     g = dest['qux']
-    assert isinstance(g, Group)
     assert 'experiment' not in g.attrs
     a = g['bar/baz']
     assert a.dtype == baz.dtype
     assert a.shape == baz.shape
     assert a.chunks == baz.chunks
-    assert a.compressor == baz.compressor
+    if hasattr(a, 'compressor') and hasattr(baz, 'compressor'):
+        assert a.compressor == baz.compressor
     assert_array_equal(a[:], baz[:])
     assert 'units' not in a.attrs
+
+    # copy group, shallow
+    dest = new_dest()
+    copy(source, dest, name='eggs', shallow=True)
+    assert 'eggs' in dest
+    eggs = dest['eggs']
+    assert 'spam' in eggs
+    a = eggs['spam']
+    assert a.dtype == spam.dtype
+    assert a.shape == spam.shape
+    assert a.chunks == spam.chunks
+    if hasattr(a, 'compressor') and hasattr(spam, 'compressor'):
+        assert a.compressor == spam.compressor
+    assert_array_equal(a[:], spam[:])
+    assert 'foo' not in eggs
+    assert 'bar' not in eggs
+
+
+def test_copy_zarr_zarr():
+    # zarr -> zarr
+    _test_copy(group, group)
+
+
+try:
+    import h5py
+    have_h5py = True
+except ImportError:
+    have_h5py = False
+
+
+def temp_h5f():
+    fn = tempfile.mktemp()
+    atexit.register(os.remove, fn)
+    h5f = h5py.File(fn, mode='w')
+    return h5f
+
+
+@pytest.mark.skipif(not have_h5py, reason='h5py not installed')
+def test_copy_h5py_zarr():
+    # h5py -> zarr
+    _test_copy(temp_h5f, group)
+
+
+@pytest.mark.skipif(not have_h5py, reason='h5py not installed')
+def test_copy_zarr_h5py():
+    # zarr -> h5py
+    _test_copy(group, temp_h5f)