Repository: opendilab/treevalue (language: python)

Issue:
[BUG] When function returns a treevalue
```python
from treevalue import FastTreeValue


def func(x):
    return FastTreeValue({
        'x': x, 'y': x ** 2,
    })


f = FastTreeValue({
    'x': func,
    'y': {
        'z': func,
    }
})
v = FastTreeValue({'x': 2, 'y': {'z': 34}})
r1 = f(v)
print(r1)
``` 

The output is 
```
<FastTreeValue 0x7f13ff6491f0>
├── 'x' --> <FastTreeValue 0x7f13ff6491c0>
│           ├── 'x' --> 2
│           └── 'y' --> 4
└── 'y' --> <FastTreeValue 0x7f13ff649220>
    └── 'z' --> <FastTreeValue 0x7f13ff649130>
                ├── 'x' --> 34
                └── 'y' --> 1156
```

But the correct output should be
```
<FastTreeValue 0x7f13ff6498b0>
├── 'x' --> <FastTreeValue 0x7f13ff649a00>
│   ├── 'x' --> 2
│   └── 'y' --> 4
└── 'y' --> <FastTreeValue 0x7f13ff6499a0>
    └── 'z' --> <FastTreeValue 0x7f13ff649910>
        ├── 'x' --> 34
        └── 'y' --> 1156
```
[BUG] When function returns a treevalue
```python
from treevalue import FastTreeValue


def func(x):
    return FastTreeValue({
        'x': x, 'y': x ** 2,
    })


f = FastTreeValue({
    'x': func,
    'y': {
        'z': func,
    }
})
v = FastTreeValue({'x': 2, 'y': {'z': 34}})
r1 = f(v)
print(r1)
``` 

The output is 
```
<FastTreeValue 0x7f13ff6491f0>
├── 'x' --> <FastTreeValue 0x7f13ff6491c0>
│           ├── 'x' --> 2
│           └── 'y' --> 4
└── 'y' --> <FastTreeValue 0x7f13ff649220>
    └── 'z' --> <FastTreeValue 0x7f13ff649130>
                ├── 'x' --> 34
                └── 'y' --> 1156
```

But the correct output should be
```
<FastTreeValue 0x7f13ff6498b0>
├── 'x' --> <FastTreeValue 0x7f13ff649a00>
│   ├── 'x' --> 2
│   └── 'y' --> 4
└── 'y' --> <FastTreeValue 0x7f13ff6499a0>
    └── 'z' --> <FastTreeValue 0x7f13ff649910>
        ├── 'x' --> 34
        └── 'y' --> 1156
```


Discussion / hints:

Fix patch:
diff --git a/treevalue/tree/func/cfunc.pyx b/treevalue/tree/func/cfunc.pyx
index 535fc77b70..33baf94a28 100644
--- 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


Test patch:
diff --git a/test/tree/func/test_func.py b/test/tree/func/test_func.py
index 1a56fe13d4..8e1b191005 100644
--- 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,
+                }
+            }
+        })