Repository: opendilab/treevalue (language: python)

Issue:
Development of flatten and unflatten for supporting parallel calculation
* [x] `flatten` - dump a tree to list of node information and value
* [x] `unflatten` - recover a list of information back to tree (inverse operation of `flatten`)




Discussion / hints:

Fix patch:
diff --git a/docs/source/api_doc/tree/tree.rst b/docs/source/api_doc/tree/tree.rst
index fbe7907b40..b86fea58bf 100644
--- a/docs/source/api_doc/tree/tree.rst
+++ b/docs/source/api_doc/tree/tree.rst
@@ -44,6 +44,22 @@ walk
 .. autofunction:: walk
 
 
+.. _apidoc_tree_tree_flatten:
+
+flatten
+-------------------
+
+.. autofunction:: flatten
+
+
+.. _apidoc_tree_tree_unflatten:
+
+unflatten
+-------------------
+
+.. autofunction:: unflatten
+
+
 .. _apidoc_tree_tree_mapping:
 
 mapping
diff --git a/docs/source/tutorials/advanced_usage/flatten_demo.demo.py b/docs/source/tutorials/advanced_usage/flatten_demo.demo.py
new file mode 100644
index 0000000000..09ecd8c341
--- /dev/null
+++ b/docs/source/tutorials/advanced_usage/flatten_demo.demo.py
@@ -0,0 +1,15 @@
+from treevalue import TreeValue, raw, flatten
+
+if __name__ == '__main__':
+    t = TreeValue({
+        'a': 1,
+        'b': 2,
+        'c': raw({'x': 3, 'y': 4}),
+        'd': {
+            'x': 3,
+            'y': 4
+        },
+    })
+
+    print('flatten(t):')
+    print(flatten(t))
diff --git a/docs/source/tutorials/advanced_usage/index.rst b/docs/source/tutorials/advanced_usage/index.rst
index 2ab64a3de8..e144deb5e5 100644
--- a/docs/source/tutorials/advanced_usage/index.rst
+++ b/docs/source/tutorials/advanced_usage/index.rst
@@ -877,6 +877,98 @@ For further informaon of function ``typetrans``, \
 take a look at :ref:`apidoc_tree_tree_typetrans`.
 
 
+Walk
+~~~~~~~~~~~~~~~
+
+You can use function :func:`treevalue.tree.tree.walk` to iterate all the nodes \
+in the tree, like the example below:
+
+.. literalinclude:: walk_demo.demo.py
+    :language: python
+    :linenos:
+
+The output should be like below:
+
+.. literalinclude:: walk_demo.demo.py.txt
+    :language: text
+    :linenos:
+
+For further informaon of function ``walk``, \
+take a look at :ref:`apidoc_tree_tree_walk`.
+
+
+Flatten Utilities
+----------------------
+
+In order to support the parallel calculation in values of \
+``TreeValue`` object, :func:`treevalue.tree.tree.flatten` and \
+:func:`treevalue.tree.tree.unflatten` is provided to dump a \
+sequence of nodes' information and value, and recover it to \
+original tree structure when calculation is completed.
+
+Flatten
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The values and node structures can be dumped \
+by function :func:`treevalue.tree.tree.flatten` as a list, \
+like this:
+
+.. literalinclude:: flatten_demo.demo.py
+    :language: python
+    :linenos:
+
+The result should be like below:
+
+.. literalinclude:: flatten_demo.demo.py.txt
+    :language: python
+    :linenos:
+
+.. note::
+
+    Function :func:`treevalue.tree.tree.flatten` is different from \
+    :func:`treevalue.tree.tree.walk` because ``flatten`` \
+    has fewer features because it is designed entirely for \
+    the dumping process before parallel calculation.
+
+    So when you need to do parallel calculation, please use \
+    :func:`treevalue.tree.tree.flatten` to dump the values to ensure the \
+    speed performance.
+
+For further informaon of function ``flatten``, \
+take a look at :ref:`apidoc_tree_tree_flatten`.
+
+
+Unflatten
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The dumped nodes' information and value can be \
+recovered to a ``TreeValue`` object by \
+function :func:`treevalue.tree.tree.unflatten`, like this:
+
+.. literalinclude:: unflatten_demo.demo.py
+    :language: python
+    :linenos:
+
+The result should be like below:
+
+.. literalinclude:: unflatten_demo.demo.py.txt
+    :language: text
+    :linenos:
+
+.. note::
+
+    It is recommended to pass an ordered iterable object \
+    in ``pairs``, this will improve the speed performance of \
+    function :func:`treevalue.tree.tree.unflatten`.
+
+    Because of this, it is a good idea to keep \
+    the :func:`treevalue.tree.tree.flatten`'s result's order \
+    when executing your own parallel processing logic.
+
+For further informaon of function ``unflatten``, \
+take a look at :ref:`apidoc_tree_tree_unflatten`.
+
+
 IO Utilities
 -----------------------
 
diff --git a/docs/source/tutorials/advanced_usage/unflatten_demo.demo.py b/docs/source/tutorials/advanced_usage/unflatten_demo.demo.py
new file mode 100644
index 0000000000..3cd016f543
--- /dev/null
+++ b/docs/source/tutorials/advanced_usage/unflatten_demo.demo.py
@@ -0,0 +1,13 @@
+from treevalue import unflatten
+
+if __name__ == '__main__':
+    flatted = [
+        (('a',), 1),
+        (('b',), 2),
+        (('c',), {'x': 3, 'y': 4}),
+        (('d', 'x'), 3),
+        (('d', 'y'), 4)
+    ]
+
+    print('unflatten(flatted):')
+    print(unflatten(flatted))
diff --git a/docs/source/tutorials/advanced_usage/walk_demo.demo.py b/docs/source/tutorials/advanced_usage/walk_demo.demo.py
new file mode 100644
index 0000000000..33d52f7e39
--- /dev/null
+++ b/docs/source/tutorials/advanced_usage/walk_demo.demo.py
@@ -0,0 +1,15 @@
+from treevalue import TreeValue, raw, walk
+
+if __name__ == '__main__':
+    t = TreeValue({
+        'a': 1,
+        'b': 2,
+        'c': raw({'x': 3, 'y': 4}),
+        'd': {
+            'x': 3,
+            'y': 4
+        },
+    })
+
+    for path, node in walk(t):
+        print(path, '-->', node)
diff --git a/treevalue/tree/tree/__init__.py b/treevalue/tree/tree/__init__.py
index f200212fd1..030d35e4f3 100644
--- a/treevalue/tree/tree/__init__.py
+++ b/treevalue/tree/tree/__init__.py
@@ -1,3 +1,4 @@
+from .flatten import flatten, unflatten
 from .functional import mapping, filter_, mask, reduce_
 from .graph import graphics
 from .io import loads, load, dumps, dump
diff --git a/treevalue/tree/tree/flatten.pxd b/treevalue/tree/tree/flatten.pxd
new file mode 100644
index 0000000000..bdaddc48be
--- /dev/null
+++ b/treevalue/tree/tree/flatten.pxd
@@ -0,0 +1,13 @@
+# distutils:language=c++
+# cython:language_level=3
+
+# flatten, unflatten
+
+from .tree cimport TreeValue
+from ..common.storage cimport TreeStorage
+
+cdef void _c_flatten(TreeStorage st, tuple path, list res) except *
+cpdef list flatten(TreeValue tree)
+
+cdef TreeStorage _c_unflatten(object pairs)
+cpdef TreeValue unflatten(object pairs, object return_type= *)
diff --git a/treevalue/tree/tree/flatten.pyx b/treevalue/tree/tree/flatten.pyx
new file mode 100644
index 0000000000..ad1001eb06
--- /dev/null
+++ b/treevalue/tree/tree/flatten.pyx
@@ -0,0 +1,105 @@
+# distutils:language=c++
+# cython:language_level=3
+
+# flatten, unflatten
+
+import cython
+
+from .tree cimport TreeValue
+from ..common.storage cimport TreeStorage
+
+cdef void _c_flatten(TreeStorage st, tuple path, list res) except *:
+    cdef dict data = st.detach()
+    cdef tuple curpath
+
+    cdef str k
+    cdef object v
+    for k, v in data.items():
+        curpath = path + (k,)
+        if isinstance(v, TreeStorage):
+            _c_flatten(v, curpath, res)
+        else:
+            res.append((curpath, v))
+
+@cython.binding(True)
+cpdef list flatten(TreeValue tree):
+    r"""
+    Overview:
+        Flatten the values in the tree.
+
+    Arguments:
+        - tree (:obj:`TreeValue`): Tree object to be flatten. 
+
+    Returns:
+        - flatted (:obj:`list`): Flatted tree, a list of tuple with (path, value).
+    
+    .. note::
+        
+        The result of function :func:`flatten` is guaranteed to be ordered, \
+        which means it obey one of the tree traversal order. But please note \
+        that the order of key values under the same subtree is not guaranteed.
+    """
+    cdef list result = []
+    _c_flatten(tree._detach(), (), result)
+    return result
+
+cdef TreeStorage _c_unflatten(object pairs):
+    cdef dict raw_data = {}
+    cdef TreeStorage result = TreeStorage(raw_data)
+    cdef list stack = []
+    stack.append(((), raw_data))
+
+    cdef tuple path
+    cdef object v
+    cdef tuple curpath, newpath
+    cdef dict curdata, newdata
+    cdef int curlen, curplen, i
+    for path, v in pairs:
+        curpath, curdata = stack[-1]
+        while path[:len(curpath)] != curpath:
+            stack.pop()
+            curpath, curdata = stack[-1]
+
+        curlen = len(curpath)
+        curplen = len(path)
+        for i in range(curlen, curplen):
+            if i < curplen - 1:
+                newpath = curpath + (path[i],)
+                if path[i] not in curdata:
+                    newdata = {}
+                    curdata[path[i]] = TreeStorage(newdata)
+                else:
+                    newdata = curdata[path[i]].detach()
+
+                curpath, curdata = newpath, newdata
+                stack.append((curpath, curdata))
+
+            else:
+                curdata[path[i]] = v
+
+    return result
+
+@cython.binding(True)
+cpdef TreeValue unflatten(object pairs, object return_type=None):
+    r"""
+    Overview:
+        Unflatten the given pairs of tree's data.
+
+    Arguments:
+        - pairs: Data pairs, should be a iterable object with items of (path, value). 
+        - return_type: Return type of unflatted tree, default is ``None`` which means use the default \
+            :class:`TreeValue` class.
+
+    Returns:
+        - tree (:obj:`TreeValue`): Unflatted tree object.
+    
+    .. note::
+        
+        It is recommended to pass an ordered iterable object in ``pairs``, this \
+        will improve the speed performance of function :func:`unflatten`.
+        
+        Because of this, it is a good idea to keep the :func:`flatten`'s result's order \
+        when executing your own processing logic.
+    """
+    return_type = return_type or TreeValue
+    return return_type(_c_unflatten(pairs))


Test patch:
diff --git a/test/tree/tree/test_flatten.py b/test/tree/tree/test_flatten.py
new file mode 100644
index 0000000000..2b6bb6ff65
--- /dev/null
+++ b/test/tree/tree/test_flatten.py
@@ -0,0 +1,60 @@
+import pytest
+
+from treevalue.tree import TreeValue, raw, flatten, unflatten
+
+
+class MyTreeValue(TreeValue):
+    pass
+
+
+# noinspection DuplicatedCode
+@pytest.mark.unittest
+class TestTreeTreeFlatten:
+    def test_flatten(self):
+        t = TreeValue({'a': 1, 'b': 2, 'c': raw({'x': 3, 'y': 4}), 'd': {'x': 3, 'y': 4}})
+
+        flatted = sorted(flatten(t))
+        assert flatted == [
+            (('a',), 1),
+            (('b',), 2),
+            (('c',), {'x': 3, 'y': 4}),
+            (('d', 'x'), 3),
+            (('d', 'y'), 4)
+        ]
+
+    def test_unflatten(self):
+        flatted = [
+            (('a',), 1),
+            (('b',), 2),
+            (('c',), {'x': 3, 'y': 4}),
+            (('d', 'x'), 3),
+            (('d', 'y'), 4)
+        ]
+        assert unflatten(flatted) == TreeValue({
+            'a': 1, 'b': 2,
+            'c': raw({'x': 3, 'y': 4}),
+            'd': {'x': 3, 'y': 4}}
+        )
+        assert unflatten(flatted, MyTreeValue) == MyTreeValue({
+            'a': 1, 'b': 2,
+            'c': raw({'x': 3, 'y': 4}),
+            'd': {'x': 3, 'y': 4}}
+        )
+
+        unordered_flatted = [
+            (('a',), 1),
+            (('d', 'x'), 3),
+            (('b',), 2),
+            (('c',), {'x': 3, 'y': 4}),
+            (('d', 'y'), 4)
+        ]
+        assert unflatten(unordered_flatted) == TreeValue({
+            'a': 1, 'b': 2,
+            'c': raw({'x': 3, 'y': 4}),
+            'd': {'x': 3, 'y': 4}}
+        )
+        assert unflatten(unordered_flatted, MyTreeValue) == MyTreeValue({
+            'a': 1, 'b': 2,
+            'c': raw({'x': 3, 'y': 4}),
+            'd': {'x': 3, 'y': 4}}
+        )