diff --git a/docs/source/api_doc/tree/tree.rst b/docs/source/api_doc/tree/tree.rst index a7e542b7a0fa3c50979ead6fcb8b2609487b046f..fbe7907b401c85e93fe343a998fc3affdc3ff80b 100644 --- a/docs/source/api_doc/tree/tree.rst +++ b/docs/source/api_doc/tree/tree.rst @@ -36,6 +36,14 @@ typetrans .. autofunction:: typetrans +.. _apidoc_tree_tree_walk: + +walk +------------------- + +.. autofunction:: walk + + .. _apidoc_tree_tree_mapping: mapping diff --git a/test/tree/tree/test_service.py b/test/tree/tree/test_service.py index d8afa4f21b1137a2297e3ae695fde8b717708eba..2745918085ee949732602e2e3b90ddc63b566f9b 100644 --- a/test/tree/tree/test_service.py +++ b/test/tree/tree/test_service.py @@ -1,6 +1,6 @@ import pytest -from treevalue.tree import jsonify, TreeValue, clone, typetrans, raw +from treevalue.tree import jsonify, TreeValue, clone, typetrans, raw, walk # noinspection DuplicatedCode @@ -78,3 +78,24 @@ class TestTreeTreeService: with pytest.raises(TypeError): typetrans(tv1, NonTreeValue) + + def test_walk(self): + class MyTreeValue(TreeValue): + pass + + tv1 = MyTreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}}) + + assert dict(walk(tv1)) == { + ('a',): 1, + ('b',): 2, + ('c', 'x',): 2, + ('c', 'y',): 3, + } + assert dict(walk(tv1, include_nodes=True)) == { + (): MyTreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}}), + ('a',): 1, + ('b',): 2, + ('c',): MyTreeValue({'x': 2, 'y': 3}), + ('c', 'x',): 2, + ('c', 'y',): 3, + } diff --git a/treevalue/tree/tree/__init__.py b/treevalue/tree/tree/__init__.py index ec4ee242babed66dfec77502239ac336d8861781..f200212fd124711b97961b3af86119aa6850efb4 100644 --- a/treevalue/tree/tree/__init__.py +++ b/treevalue/tree/tree/__init__.py @@ -1,6 +1,6 @@ from .functional import mapping, filter_, mask, reduce_ from .graph import graphics from .io import loads, load, dumps, dump -from .service import jsonify, clone, typetrans +from .service import jsonify, clone, typetrans, walk from .structural import subside, union, rise from .tree import TreeValue diff --git a/treevalue/tree/tree/service.pxd b/treevalue/tree/tree/service.pxd index a580dce91b9c3ace72304b934a86d34e05e77cec..9066f20fd1e20c602986862ade29f424e6aa6878 100644 --- a/treevalue/tree/tree/service.pxd +++ b/treevalue/tree/tree/service.pxd @@ -1,7 +1,9 @@ # distutils:language=c++ # cython:language_level=3 -# jsonify, clone, typetrans +# jsonify, clone, typetrans, walk + +from libcpp cimport bool from .tree cimport TreeValue @@ -9,3 +11,4 @@ cdef object _keep_object(object obj) cpdef object jsonify(TreeValue val) cpdef TreeValue clone(TreeValue t, object copy_value= *) cpdef TreeValue typetrans(TreeValue t, object return_type) +cpdef walk(TreeValue tree, bool include_nodes= *) diff --git a/treevalue/tree/tree/service.pyx b/treevalue/tree/tree/service.pyx index fc671e6c634a33f474f675abcf44f9f895872f0b..bd68f4af86f0a3c73410390d9a7b85f777dc8b65 100644 --- a/treevalue/tree/tree/service.pyx +++ b/treevalue/tree/tree/service.pyx @@ -1,13 +1,15 @@ # distutils:language=c++ # cython:language_level=3 -# jsonify, clone, typetrans +# jsonify, clone, typetrans, walk import copy import cython +from libcpp cimport bool from .tree cimport TreeValue +from ..common.storage cimport TreeStorage cdef object _keep_object(object obj): return obj @@ -78,3 +80,37 @@ cpdef TreeValue typetrans(TreeValue t, object return_type): raise TypeError("Tree value should be subclass of TreeValue, but {type} found.".format( type=repr(return_type.__name__) )) + +def _p_walk(TreeStorage tree, object type_, tuple path, bool include_nodes): + if include_nodes: + yield path, type_(tree) + + cdef dict data = tree.detach() + cdef str k + cdef object v + cdef tuple curpath + for k, v in data.items(): + curpath = path + (k,) + if isinstance(v, TreeStorage): + yield from _p_walk(v, type_, curpath, include_nodes) + else: + yield curpath, v + +@cython.binding(True) +cpdef walk(TreeValue tree, bool include_nodes=False): + r""" + Overview: + Walk the values and nodes in the tree. + The order of walk is not promised, if you need the ordered walking result, \ + just use function ``sorted`` at the outer side of :func:`walk`. + + Arguments: + - tree: Tree value object to be walked. + - include_nodes (:obj:`bool`): Not only the value nodes will be walked, + but the tree nodes as well. + + Returns: + - iter: Iterator to walk the given tree, contains 2 items, the 1st one is the full \ + path of the node, the 2nd one is the value. + """ + return _p_walk(tree._detach(), type(tree), (), include_nodes)