支持赋值时元组解包

Author: wukan1986

README.md

@@ -157,6 +157,14 @@ df = codegen_exec(df, _code_block_1, _code_block_2)  # 只执行，不保存代
 7. 支持`A[0]+B[1]+C[2]`，底层会转成`A+ts_delay(B,1)+ts_delay(C,2)`
 8. 支持`~A`,底层会转换成`Not(A)`
 9. `gp_`开头的函数都会返回对应的`cs_`函数。如`gp_func(A,B,C)`会替换成`cs_func(B,C)`,其中`A`用在了`groupby([date, A])`
+10. 支持`A,B,C=MACD()`元组解包，在底层会替换成
+
+   ```python
+   _x_0 = MACD()
+   A = unpack(_x_0, 0)
+   B = unpack(_x_0, 1)
+   C = unpack(_x_0, 2)
+   ```
 
 ## 下划线开头的变量
 

expr_codegen/_version.py

@@ -1 +1 @@
-__version__ = "0.9.2"
+__version__ = "0.10.0"

expr_codegen/codes.py

@@ -14,6 +14,23 @@ def __init__(self, convert_xor):
         # ^ 是异或还是乘方呢？
         self.convert_xor = convert_xor
 
+    def visit_Assign(self, node):
+        t = node.targets[0]
+        nodes = []
+        if isinstance(t, ast.Tuple):
+            for i, dim in enumerate(t.dims):
+                _v = ast.Call(
+                    func=ast.Name(id='unpack', ctx=ast.Load()),
+                    args=[node.value, ast.Constant(i)],
+                    keywords=[],
+                )
+                n = ast.Assign([dim], _v, ctx=ast.Load())
+                nodes.append(n)
+            return nodes
+
+        self.generic_visit(node)
+        return node
+
     def visit_Compare(self, node):
         assert len(node.comparators) == 1, f"不支持连续等号，请手工添加括号, {ast.unparse(node)}"
 

expr_codegen/pandas/code.py

@@ -67,19 +67,19 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                     exprs_dst.append(f"#" + '=' * 40 + func_name)
                 else:
                     va, ex, sym = kv
-                    func_code.append(f"    # {va} = {ex}\n    df[{va}] = {p.doprint(ex)}")
+                    func_code.append(f"    # {va} = {ex}\n    g[{va}] = {p.doprint(ex)}")
                     exprs_dst.append(f"{va} = {ex}")
                     if va not in syms_dst:
                         syms_out.append(va)
 
+            if len(groupbys['sort']) == 0:
+                groupbys['sort'] = f'df = df.sort_values(by=[_ASSET_, _DATE_]).reset_index(drop=True)'
             if k[0] == TS:
-                if len(groupbys['sort']) == 0:
-                    groupbys['sort'] = f'df = df.sort_values(by=[_ASSET_, _DATE_]).reset_index(drop=True)'
                 # 时序需要排序
-                func_code = [f'    df = df.sort_values(by=[_DATE_])'] + func_code
-            elif k[0] == CS:
-                if len(groupbys['sort']) == 0:
-                    groupbys['sort'] = f'df = df.sort_values(by=[_DATE_, _ASSET_]).reset_index(drop=True)'
+                func_code = [f'    g.df = df.sort_values(by=[_DATE_])'] + func_code
+            else:
+                # 时序需要排序
+                func_code = [f'    g.df = df'] + func_code
 
             # polars风格代码列表
             funcs[func_name] = '\n'.join(func_code)

expr_codegen/pandas/helper.py

@@ -0,0 +1,35 @@
+"""
+A、B、C=MACD()无法生成DAG，所以变通的改成
+
+A=unpack(MACD(),0)
+B=unpack(MACD(),1)
+C=unpack(MACD(),2)
+
+cse能自动提取成
+
+_x_0 = MACD()
+
+但 df['_x_0'] 是无法放入tuple的，所以决定用另一个类来实现兼容
+
+"""
+import pandas as pd
+
+
+class GlobalVariable(object):
+    def __init__(self):
+        self.dict = {}
+        self.df = pd.DataFrame()
+
+    def __getitem__(self, item):
+        if item in self.dict:
+            return self.dict[item]
+        return self.df[item]
+
+    def __setitem__(self, key, value):
+        if isinstance(value, tuple):
+            # tuple存字典中
+            self.dict[key] = value
+            # 占位，避免drop时报错
+            self.df[key] = False
+        else:
+            self.df[key] = value

expr_codegen/pandas/printer.py

@@ -54,7 +54,7 @@ def _print(self, expr, **kwargs) -> str:
             self._print_level -= 1
 
     def _print_Symbol(self, expr):
-        return f"df[{expr.name}]"
+        return f"g[{expr.name}]"
 
     def _print_Equality(self, expr):
         PREC = precedence(expr)

expr_codegen/pandas/ta.py

@@ -7,9 +7,16 @@
 
 所以有必要使用类似于polars_ta的公共库，但因目前未找到合适库，所以以下是临时版，以后要独立出去
 """
+from typing import Tuple
+
 import numpy as np
 import pandas as pd
 
+try:
+    import talib
+except:
+    pass
+
 
 def abs_(x: pd.Series) -> pd.Series:
     return np.abs(x)
@@ -93,3 +100,7 @@ def ts_std_dev(x: pd.Series, d: int = 5, ddof: int = 0) -> pd.Series:
 
 def ts_sum(x: pd.Series, d: int = 5) -> pd.Series:
     return x.rolling(d).sum()
+
+
+def ts_MACD(close: pd.Series, fastperiod: int = 12, slowperiod: int = 26, signalperiod: int = 9) -> Tuple[pd.Series, pd.Series, pd.Series]:
+    return talib.MACD(close, fastperiod, slowperiod, signalperiod)

expr_codegen/pandas/template.py.j2

@@ -6,9 +6,9 @@ import numpy as np  # noqa
 import pandas as pd  # noqa
 from loguru import logger  # noqa
 
+from expr_codegen.pandas.helper import GlobalVariable
 from expr_codegen.pandas.ta import *  # noqa
 
-
 {{ syms1 }}
 
 {{ syms2 }}
@@ -19,15 +19,19 @@ _NONE_ = None
 _TRUE_ = True
 _FALSE_ = False
 
+g = GlobalVariable()
+
+def unpack(x: Tuple, idx: int = 0) -> pd.Series:
+    return x[idx]
+
 {%-for row in extra_codes %}
 {{ row-}}
 {% endfor %}
 
-
 {% for key, value in funcs.items() %}
 def {{ key }}(df: pd.DataFrame) -> pd.DataFrame:
 {{ value }}
-    return df
+    return g.df
 {% endfor %}
 
 """

expr_codegen/polars_group/template.py.j2

@@ -30,6 +30,9 @@ _NONE_ = None
 _TRUE_ = True
 _FALSE_ = False
 
+def unpack(x: Expr, idx: int = 0) -> Expr:
+    return x.struct[idx]
+
 {%-for row in extra_codes %}
 {{ row-}}
 {% endfor %}

expr_codegen/polars_over/template.py.j2

@@ -30,6 +30,9 @@ _NONE_ = None
 _TRUE_ = True
 _FALSE_ = False
 
+def unpack(x: Expr, idx: int = 0) -> Expr:
+    return x.struct[idx]
+
 {%-for row in extra_codes %}
 {{ row-}}
 {% endfor %}