Merge pull request #58 from opendilab/dev/setdefault

dev(hansbug): add setdefault method to TreeValue

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, popitem, clear, update
+    :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, clear, update, setdefault
 
 
 .. _apidoc_tree_tree_delayed:

test/tree/common/test_storage.py

@@ -264,6 +264,26 @@ class TestTreeStorage:
         t.set('fff', raw({'x': 1, 'y': 2}))
         assert t.get('fff') == {'x': 1, 'y': 2}
 
+    def test_setdefault(self):
+        t = create_storage({})
+        assert t.setdefault('a', 1) == 1
+        assert t == create_storage({'a': 1})
+        assert t.setdefault('b', 2) == 2
+        assert t == create_storage({'a': 1, 'b': 2})
+        assert t.setdefault('a', 100) == 1
+        assert t == create_storage({'a': 1, 'b': 2})
+
+        assert t.setdefault('c', create_storage({'a': 100, 'b': 200})) == create_storage({'a': 100, 'b': 200})
+        assert t == create_storage({'a': 1, 'b': 2, 'c': {'a': 100, 'b': 200}})
+        assert t.setdefault('c', create_storage({'a': 400, 'b': 300})) == create_storage({'a': 100, 'b': 200})
+        assert t == create_storage({'a': 1, 'b': 2, 'c': {'a': 100, 'b': 200}})
+
+        d1 = delayed_partial(lambda: 1)
+        assert t.setdefault('g', delayed_partial(lambda x: x + 1, d1)) == 2
+        assert t == create_storage({'a': 1, 'b': 2, 'c': {'a': 100, 'b': 200}, 'g': 2})
+        assert t.setdefault('g', delayed_partial(lambda x: x * 100, d1)) == 2
+        assert t == create_storage({'a': 1, 'b': 2, 'c': {'a': 100, 'b': 200}, 'g': 2})
+
     def test_del_(self):
         t = create_storage({'a': 1, 'b': 2, 'c': raw({'x': 3, 'y': 4}), 'd': {'x': 3, 'y': 4}})
         t.del_('c')

test/tree/tree/test_tree.py

@@ -146,6 +146,29 @@ class TestTreeTreeTree:
         assert isinstance(tv1.fff, dict)
         assert tv1.fff == {'x': 1, 'y': 2}
 
+    def test_setdefault(self):
+        t = TreeValue({})
+        assert t.setdefault('a', 1) == 1
+        assert t == TreeValue({'a': 1})
+        assert t.setdefault('a', 100) == 1
+        assert t == TreeValue({'a': 1})
+
+        assert t.setdefault('f', {'a': 1, 'b': 2}) == {'a': 1, 'b': 2}
+        assert t == TreeValue({'a': 1, 'f': raw({'a': 1, 'b': 2})})
+        assert t.setdefault('f', {'y': 1, 'z': 2}) == {'a': 1, 'b': 2}
+        assert t == TreeValue({'a': 1, 'f': raw({'a': 1, 'b': 2})})
+
+        assert t.setdefault('c', TreeValue({'a': 1, 'b': 2})) == TreeValue({'a': 1, 'b': 2})
+        assert t == TreeValue({'a': 1, 'f': raw({'a': 1, 'b': 2}), 'c': {'a': 1, 'b': 2}})
+        assert t.setdefault('c', TreeValue({'aa': 1, 'bb': 2})) == TreeValue({'a': 1, 'b': 2})
+        assert t == TreeValue({'a': 1, 'f': raw({'a': 1, 'b': 2}), 'c': {'a': 1, 'b': 2}})
+
+        d = delayed(lambda: 1)
+        assert t.setdefault('g', delayed(lambda x: x + 1, d)) == 2
+        assert t == TreeValue({'a': 1, 'f': raw({'a': 1, 'b': 2}), 'c': {'a': 1, 'b': 2}, 'g': 2})
+        assert t.setdefault('g', delayed(lambda x: x + 100, d)) == 2
+        assert t == TreeValue({'a': 1, 'f': raw({'a': 1, 'b': 2}), 'c': {'a': 1, 'b': 2}, 'g': 2})
+
     def test_tree_value_operate_with_item(self):
         tv1 = TreeValue({'a': 1, 'b': 2, 'c': {'x': 2, 'y': 3}})
         tv2 = TreeValue(tv1)

treevalue/tree/common/storage.pxd

@@ -10,6 +10,7 @@ cdef class TreeStorage:
     cdef readonly dict map
 
     cpdef public void set(self, str key, object value) except *
+    cpdef public object setdefault(self, str key, object default)
     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)

treevalue/tree/common/storage.pyx

@@ -22,6 +22,16 @@ cdef class TreeStorage:
     cpdef public void set(self, str key, object value) except *:
         self.map[key] = unraw(value)
 
+    cpdef public object setdefault(self, str key, object default):
+        cdef object v, df
+        try:
+            v = self.map[key]
+            return _c_undelay_data(self.map, key, v)
+        except KeyError:
+            df = unraw(default)
+            self.map[key] = df
+            return _c_undelay_data(self.map, key, df)
+
     # 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

treevalue/tree/tree/tree.pxd

@@ -25,6 +25,7 @@ cdef class TreeValue:
     cpdef public pop(self, str key, object default= *)
     cpdef public popitem(self)
     cpdef public void clear(self)
+    cpdef public object setdefault(self, str key, object default= *)
 
     cpdef public treevalue_keys keys(self)
     cpdef public treevalue_values values(self)

treevalue/tree/tree/tree.pyx

@@ -169,7 +169,66 @@ cdef class TreeValue:
         """
         self._st.clear()
 
-    cdef void _update(self, object d, dict kwargs) except *:
+    @cython.binding(True)
+    cpdef object setdefault(self, str key, object default=None):
+        """
+        Overview:
+            Set the ``default`` to this treevalue and return it if the ``key`` is not exist, \
+            otherwise just return the existing value of ``key``.
+
+        :param key: Items' name.
+        :param default: Default value of the ``key``, ``None`` will be used when not given.
+        :return: The newest value of the ``key``.
+
+        .. note::
+            The behaviour of method :meth:`setdefault` is similar to \
+            `dict.setdefault <https://docs.python.org/3/library/stdtypes.html#dict.setdefault>`_.
+
+        Examples::
+            >>> from treevalue import TreeValue, delayed
+            >>> t = TreeValue({'a': 1, 'b': 3, 'c': '233'})
+            >>> t.setdefault('d', 'dsflgj')  # set new value
+            'dsflgj'
+            >>> t
+            <TreeValue 0x7efe31576048>
+            ├── 'a' --> 1
+            ├── 'b' --> 3
+            ├── 'c' --> '233'
+            └── 'd' --> 'dsflgj'
+            >>>
+            >>> t.setdefault('ff')  # default value - None
+            >>> t
+            <TreeValue 0x7efe31576048>
+            ├── 'a' --> 1
+            ├── 'b' --> 3
+            ├── 'c' --> '233'
+            ├── 'd' --> 'dsflgj'
+            └── 'ff' --> None
+            >>>
+            >>> t.setdefault('a', 1000)  # existing key
+            1
+            >>> t
+            <TreeValue 0x7efe31576048>
+            ├── 'a' --> 1
+            ├── 'b' --> 3
+            ├── 'c' --> '233'
+            ├── 'd' --> 'dsflgj'
+            └── 'ff' --> None
+            >>>
+            >>> t.setdefault('g', delayed(lambda: 1))  # delayed value
+            1
+            >>> t
+            <TreeValue 0x7efe31576048>
+            ├── 'a' --> 1
+            ├── 'b' --> 3
+            ├── 'c' --> '233'
+            ├── 'd' --> 'dsflgj'
+            ├── 'ff' --> None
+            └── 'g' --> 1
+        """
+        return self._unraw(self._st.setdefault(key, self._raw(default)))
+
+    cdef inline void _update(self, object d, dict kwargs) except *:
         cdef object dt
         if d is None:
             dt = {}