fix(hansbug): fix bug of #82, add more unittests

opendilab · Mar 6, 2023 · 80d4e9b · 80d4e9b
1 parent 15debc5
commit 80d4e9b
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git a/test/tree/func/test_func.py b/test/tree/func/test_func.py
@@ -3,6 +3,7 @@
 
 import pytest
 
+from treevalue import FastTreeValue
 from treevalue.tree import func_treelize, TreeValue, method_treelize, classmethod_treelize, delayed
 
 
@@ -401,3 +402,29 @@ def total(a, b):
             'v': {'a': 12, 'b': 25, 'x': {'c': 38, 'd': 51}},
         })
         assert cnt_1 == 4
+
+    def test_return_treevalue(self):
+        def func(x):
+            return FastTreeValue({
+                'x': x, 'y': x ** 2,
+            })
+
+        f = FastTreeValue({
+            'x': func,
+            'y': {
+                'z': func,
+            }
+        })
+        v = FastTreeValue({'x': 2, 'y': {'z': 34}})
+        assert f(v) == FastTreeValue({
+            'x': {
+                'x': v.x,
+                'y': v.x ** 2,
+            },
+            'y': {
+                'z': {
+                    'x': v.y.z,
+                    'y': v.y.z ** 2,
+                }
+            }
+        })
diff --git a/treevalue/tree/func/cfunc.pyx b/treevalue/tree/func/cfunc.pyx
@@ -57,6 +57,7 @@ cdef object _c_func_treelize_run(object func, list args, dict kwargs, _e_tree_mo
 
     cdef list _a_args
     cdef dict _a_kwargs
+    cdef object _a_ret
     if not has_tree:
         _a_args = []
         for v in args:
@@ -72,7 +73,11 @@ cdef object _c_func_treelize_run(object func, list args, dict kwargs, _e_tree_mo
             else:
                 _a_kwargs[k] = missing_func()
 
-        return func(*_a_args, **_a_kwargs)
+        _a_ret = func(*_a_args, **_a_kwargs)
+        if isinstance(_a_ret, TreeValue):
+            return _a_ret._detach()
+        else:
+            return _a_ret
 
     cdef dict _d_res = {}
     cdef str ak