Merge pull request #37 from opendilab/dev/pop

dev(hansbug): add pop method for TreeValue

Author: HansBug

docs/source/api_doc/tree/common.rst

@@ -9,7 +9,7 @@ TreeStorage
 -----------
 
 .. autoclass:: TreeStorage
-    :members: get, set, del_, contains, size, empty, copy, deepcopy, deepcopyx, dump, deepdump, deepdumpx, jsondumpx, copy_from, deepcopy_from, deepcopyx_from, detach
+    :members: get, get_or_default, pop, pop_or_default, set, del_, contains, size, empty, copy, deepcopy, deepcopyx, dump, deepdump, deepdumpx, jsondumpx, copy_from, deepcopy_from, deepcopyx_from, detach
 
 
 .. _apidoc_tree_common_raw:

docs/source/api_doc/tree/tree.rst

@@ -9,7 +9,7 @@ TreeValue
 ---------------
 
 .. autoclass:: TreeValue
-    :members: __init__, __getattribute__, __setattr__, __delattr__, __contains__, __repr__, __iter__, __hash__, __eq__, _attr_extern, __len__, __bool__, __str__, __getstate__, __setstate__, get, keys, values, items, __getitem__, __setitem__, __delitem__, _getitem_extern, _setitem_extern, _delitem_extern
+    :members: __init__, __getattribute__, __setattr__, __delattr__, __contains__, __repr__, __iter__, __hash__, __eq__, _attr_extern, __len__, __bool__, __str__, __getstate__, __setstate__, get, pop, keys, values, items, __getitem__, __setitem__, __delitem__, _getitem_extern, _setitem_extern, _delitem_extern
 
 
 .. _apidoc_tree_tree_delayed:

test/tree/common/test_storage.py

@@ -97,6 +97,110 @@ def test_get_or_default(self):
         assert t1.get_or_default('fff', delayed_partial(lambda: 2345)) == 2345
         assert not t1.contains('fff')
 
+    def test_pop(self):
+        t = create_storage({'a': 1, 'b': 2, 'c': raw({'x': 3, 'y': 4}), 'd': {'x': 3, 'y': 4}})
+        assert t.pop('a') == 1
+        with pytest.raises(KeyError):
+            t.pop('a')
+
+        assert t.pop('b') == 2
+        assert t.pop('c') == {'x': 3, 'y': 4}
+
+        td = t.pop('d')
+        assert isinstance(td, TreeStorage)
+        assert td.get('x') == 3
+        assert td.get('y') == 4
+
+        with pytest.raises(KeyError):
+            t.pop('aksjdlasdkjf')
+
+        cnt1, cnt2, cnt3 = 0, 0, 0
+
+        def f1():
+            nonlocal cnt1
+            cnt1 += 1
+            return 2
+
+        def f2(x, y):
+            nonlocal cnt2
+            cnt2 += 1
+            return {'x': x, 'y': y}
+
+        def f3(x, y):
+            nonlocal cnt3
+            cnt3 += 1
+            return create_storage({'x': x, 'y': raw(y)})
+
+        t2 = create_storage({
+            'a': 1,
+            'b': delayed_partial(f1),
+            'c': delayed_partial(f2, delayed_partial(f1), 3),
+            'd': delayed_partial(f3, 3, delayed_partial(f2, 3, 4))
+        })
+
+        assert t2.pop('a') == 1
+
+        assert cnt1 == 0
+        assert t2.pop('b') == 2
+        assert cnt1 == 1
+        with pytest.raises(KeyError):
+            t2.pop('b')
+        assert cnt1 == 1
+
+        assert (cnt1, cnt2) == (1, 0)
+        assert t2.pop('c') == {'x': 2, 'y': 3}
+        assert (cnt1, cnt2) == (2, 1)
+        with pytest.raises(KeyError):
+            t2.pop('c')
+        assert (cnt1, cnt2) == (2, 1)
+
+        assert (cnt1, cnt2, cnt3) == (2, 1, 0)
+        assert t2.get('d').pop('x') == 3
+        assert t2.get('d').pop('y') == {'x': 3, 'y': 4}
+        assert (cnt1, cnt2, cnt3) == (2, 2, 1)
+        with pytest.raises(KeyError):
+            t2.get('d').pop('x')
+        with pytest.raises(KeyError):
+            t2.get('d').pop('y')
+        assert (cnt1, cnt2, cnt3) == (2, 2, 1)
+
+    def test_pop_or_default(self):
+        t = create_storage({'a': 1, 'b': 2, 'c': raw({'x': 3, 'y': 4}), 'd': {'x': 3, 'y': 4}})
+        assert t.pop_or_default('a', 233) == 1
+        with pytest.raises(KeyError):
+            t.pop('a')
+        assert t.pop_or_default('a', 233) == 233
+
+        assert t.pop_or_default('b', 233) == 2
+        assert t.pop_or_default('c', 233) == {'x': 3, 'y': 4}
+
+        td = t.pop_or_default('d', 233)
+        assert isinstance(td, TreeStorage)
+        assert td.pop_or_default('x', 233) == 3
+        assert td.pop_or_default('y', 233) == 4
+
+        assert t.pop_or_default('fff', 233) == 233
+
+        t = create_storage({'a': 1, 'b': 2, 'c': raw({'x': 3, 'y': 4}), 'd': {'x': 3, 'y': 4}})
+        t1 = create_storage({
+            'a': delayed_partial(lambda: t.get('a')),
+            'b': delayed_partial(lambda: t.get('b')),
+            'c': delayed_partial(lambda: t.get('c')),
+            'd': delayed_partial(lambda: t.get('d')),
+        })
+        assert t1.pop_or_default('a', 233) == 1
+        assert t1.pop_or_default('b', 233) == 2
+        assert t1.pop_or_default('c', 233) == {'x': 3, 'y': 4}
+
+        t1d = t1.pop_or_default('d', 233)
+        assert isinstance(t1d, TreeStorage)
+        assert t1d.pop_or_default('x', 233) == 3
+        assert t1d.pop_or_default('y', 233) == 4
+
+        assert t1.pop_or_default('fff', 233) == 233
+        assert t1.pop_or_default('fff', delayed_partial(lambda: 2345)) == 2345
+        assert not t1.contains('fff')
+
     def test_set(self):
         t = create_storage({})
         t.set('a', 1)

test/tree/tree/test_tree.py

@@ -282,6 +282,52 @@ def test_get(self):
             _ = tv1.get('e')
         assert tv1.get('e', 233) == 233
 
+        tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}, 'd': raw({'x': 2, 'y': 3})})
+        tv2 = TreeValue({
+            'a': delayed(lambda: tv1.a),
+            'b': delayed(lambda: tv1.b),
+            'c': delayed(lambda: tv1.c),
+        })
+        assert tv2.get('a') == 1
+        assert tv2.get('b') == 2
+        assert tv2.get('c') == TreeValue({'x': 2, 'y': 3})
+
+    def test_pop(self):
+        tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}, 'd': raw({'x': 2, 'y': 3})})
+
+        assert tv1.pop('a') == 1
+        with pytest.raises(KeyError):
+            _ = tv1.pop('a')
+        assert tv1.pop('a', 233) == 233
+
+        assert tv1.pop('b') == 2
+        assert tv1.pop('c') == TreeValue({'x': 2, 'y': 3})
+        assert tv1.pop('d') == {'x': 2, 'y': 3}
+        with pytest.raises(KeyError):
+            _ = tv1.pop('e')
+        assert tv1.pop('e', 233) == 233
+
+        tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}, 'd': raw({'x': 2, 'y': 3})})
+        tv2 = TreeValue({
+            'a': delayed(lambda: tv1.a),
+            'b': delayed(lambda: tv1.b),
+            'c': delayed(lambda: tv1.c),
+        })
+        assert tv2.pop('a') == 1
+        assert tv2.pop('b') == 2
+        assert tv2.pop('c') == TreeValue({'x': 2, 'y': 3})
+
+        tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}, 'd': raw({'x': 2, 'y': 3})})
+        tv3 = TreeValue({'a': 233, 'b': tv1, 'c': tv1})
+        assert tv3.pop('b') == tv1
+        assert tv3.pop('c') == tv1
+        with pytest.raises(KeyError):
+            tv3.pop('b')
+        with pytest.raises(KeyError):
+            tv3.pop('c')
+        assert tv3.pop('b', 345) == 345
+        assert tv3.pop('c', 345) == 345
+
     def test_keys(self):
         tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}, 'd': raw({'x': 2, 'y': 3})})
         assert set(tv1.keys()) == {'a', 'b', 'c', 'd'}

treevalue/tree/common/storage.pxd

@@ -12,6 +12,8 @@ cdef class TreeStorage:
     cpdef public void set(self, str key, object value) except *
     cpdef public object get(self, str key)
     cpdef public object get_or_default(self, str key, object default)
+    cpdef public object pop(self, str key)
+    cpdef public object pop_or_default(self, str key, object default)
     cpdef public void del_(self, str key) except *
     cpdef public boolean contains(self, str key)
     cpdef public uint size(self)

treevalue/tree/common/storage.pyx

@@ -22,6 +22,7 @@ cdef class TreeStorage:
     cpdef public void set(self, str key, object value) except *:
         self.map[key] = value
 
+    # get and get_or_default is designed separately due to the consideration of performance
     cpdef public object get(self, str key):
         cdef object v, nv
         try:
@@ -35,6 +36,25 @@ cdef class TreeStorage:
         v = self.map.get(key, default)
         return _c_undelay_check_data(self.map, key, v)
 
+    # pop and pop_or_default is designed separately due to the consideration of performance
+    cpdef public object pop(self, str key):
+        cdef object v, nv, res
+        try:
+            v = self.map[key]
+            res = _c_undelay_data(self.map, key, v)
+            del self.map[key]
+            return res
+        except KeyError:
+            raise KeyError(f"Key {repr(key)} not found in this tree.")
+
+    cpdef public object pop_or_default(self, str key, object default):
+        cdef object v, nv, res
+        v = self.map.get(key, default)
+        res = _c_undelay_data(self.map, key, v)
+        if key in self.map:
+            del self.map[key]
+        return res
+
     cpdef public void del_(self, str key) except *:
         try:
             del self.map[key]

treevalue/tree/tree/tree.pxd

@@ -18,6 +18,7 @@ cdef class TreeValue:
     cpdef _setitem_extern(self, object key, object value)
     cpdef _delitem_extern(self, object key)
     cpdef get(self, str key, object default= *)
+    cpdef pop(self, str key, object default= *)
 
 cdef str _prefix_fix(object text, object prefix)
 cdef object _build_tree(TreeStorage st, object type_, str prefix, dict id_pool, tuple path)
@@ -28,4 +29,4 @@ cdef class DetachedDelayedProxy(DelayedProxy):
     cdef object val
 
     cpdef object value(self)
-    cpdef object fvalue(self)
+    cpdef object fvalue(self)

treevalue/tree/tree/tree.pyx

@@ -119,6 +119,28 @@ cdef class TreeValue:
 
         return self._unraw(value)
 
+    @cython.binding(True)
+    cpdef pop(self, str key, object default=_GET_NO_DEFAULT):
+        """
+        Overview:
+            Pop item from the tree node.
+
+        Arguments:
+            - key (:obj:`str`): Item's name.
+            - default (:obj:`default`): Default value when this item is not found, default is \
+                `_GET_NO_DEFAULT` which means just raise `KeyError` when not found.
+
+        Returns:
+            - value: Item's value.
+        """
+        cdef object value
+        if default is _GET_NO_DEFAULT:
+            value = self._st.pop(key)
+        else:
+            value = self._st.pop_or_default(key, default)
+
+        return self._unraw(value)
+
     @cython.binding(True)
     cpdef _attr_extern(self, str key):
         r"""