Merge pull request #52 from opendilab/dev/mutablemapping

dev(hansbug): add popitem method for TreeValue

Author: HansBug

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, pop, 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, popitem
 
 
 .. _apidoc_tree_tree_delayed:

test/tree/common/test_storage.py

@@ -208,6 +208,33 @@ def test_pop_or_default(self):
         assert t1.pop_or_default('fff', delayed_partial(lambda: 2345)) == 2345
         assert not t1.contains('fff')
 
+    def test_popitem(self):
+        t = create_storage({})
+        with pytest.raises(KeyError):
+            t.popitem()
+
+        t1 = create_storage({'a': 1, 'b': 2, 'c': {'x': 2}})
+        assert sorted([t1.popitem() for _ in range(t1.size())]) == [
+            ('a', 1), ('b', 2), ('c', create_storage({'x': 2}))
+        ]
+
+        t2 = create_storage({
+            'a': delayed_partial(lambda: 1),
+            'b': delayed_partial(lambda: 2),
+        })
+        assert sorted([t2.popitem() for _ in range(t2.size())]) == [
+            ('a', 1), ('b', 2)
+        ]
+
+        d1 = delayed_partial(lambda: 1)
+        d2 = delayed_partial(lambda x: x + 1, d1)
+        t3 = create_storage({
+            'a': d1, 'b': d2, 'c': d1,
+        })
+        assert sorted([t3.popitem() for _ in range(t3.size())]) == [
+            ('a', 1), ('b', 2), ('c', 1)
+        ]
+
     def test_set(self):
         t = create_storage({})
         t.set('a', 1)

test/tree/tree/test_tree.py

@@ -336,6 +336,25 @@ def test_pop(self):
         assert tv3.pop('b', 345) == 345
         assert tv3.pop('c', 345) == 345
 
+    def test_popitem(self):
+        tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}, 'd': raw({'x': 2, 'y': 3})})
+        assert sorted([tv1.popitem() for _ in range(len(tv1))]) == [
+            ('a', 1), ('b', 2),
+            ('c', TreeValue({'x': 2, 'y': 3})),
+            ('d', {'x': 2, 'y': 3}),
+        ]
+        with pytest.raises(KeyError):
+            tv1.popitem()
+
+        d1 = delayed(lambda: 1)
+        d2 = delayed(lambda x: x + 1, d1)
+        tv2 = TreeValue({'a': d1, 'b': d2, 'c': d1, 'd': 100})
+        assert sorted([tv2.popitem() for _ in range(len(tv2))]) == [
+            ('a', 1), ('b', 2), ('c', 1), ('d', 100),
+        ]
+        with pytest.raises(KeyError):
+            tv2.popitem()
+
     def test_keys(self):
         tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}, 'd': raw({'x': 2, 'y': 3})})
         assert len(tv1.keys()) == 4

treevalue/tree/common/storage.pxd

@@ -14,6 +14,7 @@ cdef class TreeStorage:
     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 tuple popitem(self)
     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

@@ -55,6 +55,12 @@ cdef class TreeStorage:
             del self.map[key]
         return res
 
+    cpdef public tuple popitem(self):
+        cdef str k
+        cdef object v
+        k, v = self.map.popitem()
+        return k, undelay(v)
+
     cpdef public void del_(self, str key) except *:
         try:
             del self.map[key]

treevalue/tree/tree/tree.pxd

@@ -22,6 +22,7 @@ cdef class TreeValue:
     cpdef _delitem_extern(self, object key)
     cpdef get(self, str key, object default= *)
     cpdef pop(self, str key, object default= *)
+    cpdef popitem(self)
 
     cpdef treevalue_keys keys(self)
     cpdef treevalue_values values(self)

treevalue/tree/tree/tree.pyx

@@ -140,6 +140,27 @@ cdef class TreeValue:
 
         return self._unraw(value)
 
+    @cython.binding(True)
+    cpdef popitem(self):
+        """
+        Overview:
+            Pop item (with a key and its value) from the tree node.
+
+        :return: Popped item.
+        :raise KeyError: When current treevalue is empty.
+
+        .. note::
+            The method :meth:`popitem` will raise ``KeyError`` when empty, like the behaviour in \
+            `dict.popitem <https://docs.python.org/3/library/stdtypes.html#dict.popitem>`_.
+        """
+        cdef str k
+        cdef object v
+        try:
+            k, v = self._st.popitem()
+            return k, self._unraw(v)
+        except KeyError:
+            raise KeyError(f'popitem(): {self._type.__name__} is empty.')
+
     @cython.binding(True)
     cpdef _attr_extern(self, str key):
         r"""