expr-codegen 0.10.15__tar.gz → 0.11.0__tar.gz

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: expr_codegen
-Version: 0.10.15
+Version: 0.11.0
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License
@@ -43,6 +43,7 @@ Requires-Dist: Jinja2
 Requires-Dist: networkx
 Requires-Dist: loguru
 Requires-Dist: sympy
+Requires-Dist: ast-comments
 Provides-Extra: streamlit
 Requires-Dist: streamlit; extra == "streamlit"
 Requires-Dist: streamlit-ace; extra == "streamlit"
@@ -211,6 +212,33 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 2. `over_null='order_by'`。分到一个区域，`null`排在前面
 3. `over_null=None`。不处理，直接调用，速度更快。如果确信不会中段产生`null`建议使用此参数
 
+`codegen_exec(over_null='partition_by')`为全局使用`partition_by`。但遇到`ts_count_nulls`这类`null`
+函数就得使用`over_null=None`，所以本工具还新添了注释功能来指定单行表达式参数
+
+1. `# --over_null partition_by`。单行`over_null='partition_by'`
+2. `# --over_null=order_by`。单行`over_null='order_by'`
+3. `# --over_null`。单行`over_null=None`
+4. `# `。取`codegen_exec`参数传入的`over_null`值
+
+注意：
+
+1. `# --over_null`传参注释只能写在单行表达式的后面，不能独立成一行，否则会被忽略
+2. `# --over_null # --over_null=order_by`多个`#`时，只取第一个有效
+3. 只对最外层`ts`函数有效。如果`ts`函数不在外层，需要人工提炼。例如：
+   ```python
+   X1 = cs_rank(ts_mean(CLOSE, 3)) # --over_null=order_by # 应用在cs_rank上，没有意义
+   X2 = ts_rank(ts_mean(CLOSE, 3), 5) # --over_null=order_by # 本以为应用在ts_rank(ts_mean)上，但由于出现了公共ts_mean，其实是应用在ts_rank(_x_0)上
+   ```
+
+   需写成
+
+   ```python
+   _x_0 = ts_mean(CLOSE, 3)  # --over_null=order_by 
+   X1 = cs_rank(_x_0)
+   X2 = ts_rank(_x_0, 5)
+   ```
+4. 由于很容易搞错，强烈建议生成`output_file`,检查生成的代码是否正确。
+
 ## `expr_codegen`局限性
 
 1. `DAG`只能增加列无法删除。增加列时，遇到同名列会覆盖
@@ -220,7 +248,8 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 
 ## 特别语法
 
-1. 支持`C?T:F`三元表达式（仅可字符串中使用），底层会先转成`C or True if( T )else F`，然后修正成`T if C else F`，最后转成`if_else(C,T,F)`。支持与`if else`混用
+1. 支持`C?T:F`三元表达式（仅可字符串中使用），底层会先转成`C or True if( T )else F`，然后修正成`T if C else F`
+   ，最后转成`if_else(C,T,F)`。支持与`if else`混用
 2. `(A<B)*-1`,底层将转换成`int_(A<B)*-1`
 3. 为防止`A==B`被`sympy`替换成`False`，底层会换成`Eq(A,B)`
 4. `A^B`的含义与`convert_xor`参数有关，`convert_xor=True`底层会转换成`Pow(A,B)`，反之为`Xor(A,B)`。默认为`False`，用`**`表示乘方
@@ -230,6 +259,7 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 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()`元组解包，在底层会替换成
+11. 单行注释支持参数输入，如：`# --over_null`、`# --over_null=order_by`、`# --over_null=partition_by`
 
 ```python
 _x_0 = MACD()
@@ -242,7 +272,8 @@ C = unpack(_x_0, 2)
 
 1. 输出的数据，所有以`_`开头的列，最后会被自动删除。所以需要保留的变量一定不要以`_`开头
 2. 为减少重复计算，自动添加了了中间变量，以`_x_`开头，如`_x_0`，`_x_1`等。最后会被自动删除
-3. 单行表达式过长，或有重复计算，可以通过中间变量，将单行表达式改成多行。如果中间变量使用`_`开头，将会自动添加数字后缀，形成不同的变量，如`_A`会替换成`_A_0_`、`_A_1_`等。使用场景如下：
+3. 单行表达式过长，或有重复计算，可以通过中间变量，将单行表达式改成多行。如果中间变量使用`_`
+   开头，将会自动添加数字后缀，形成不同的变量，如`_A`会替换成`_A_0_`、`_A_1_`等。使用场景如下：
     1. 同一变量名，重复使用。本质是不同的变量
     2. 循环赋值，但`DAG`不支持有环。`=`号左右的同名变量其实是不同变量
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/README.md

@@ -160,6 +160,33 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 2. `over_null='order_by'`。分到一个区域，`null`排在前面
 3. `over_null=None`。不处理，直接调用，速度更快。如果确信不会中段产生`null`建议使用此参数
 
+`codegen_exec(over_null='partition_by')`为全局使用`partition_by`。但遇到`ts_count_nulls`这类`null`
+函数就得使用`over_null=None`，所以本工具还新添了注释功能来指定单行表达式参数
+
+1. `# --over_null partition_by`。单行`over_null='partition_by'`
+2. `# --over_null=order_by`。单行`over_null='order_by'`
+3. `# --over_null`。单行`over_null=None`
+4. `# `。取`codegen_exec`参数传入的`over_null`值
+
+注意：
+
+1. `# --over_null`传参注释只能写在单行表达式的后面，不能独立成一行，否则会被忽略
+2. `# --over_null # --over_null=order_by`多个`#`时，只取第一个有效
+3. 只对最外层`ts`函数有效。如果`ts`函数不在外层，需要人工提炼。例如：
+   ```python
+   X1 = cs_rank(ts_mean(CLOSE, 3)) # --over_null=order_by # 应用在cs_rank上，没有意义
+   X2 = ts_rank(ts_mean(CLOSE, 3), 5) # --over_null=order_by # 本以为应用在ts_rank(ts_mean)上，但由于出现了公共ts_mean，其实是应用在ts_rank(_x_0)上
+   ```
+
+   需写成
+
+   ```python
+   _x_0 = ts_mean(CLOSE, 3)  # --over_null=order_by 
+   X1 = cs_rank(_x_0)
+   X2 = ts_rank(_x_0, 5)
+   ```
+4. 由于很容易搞错，强烈建议生成`output_file`,检查生成的代码是否正确。
+
 ## `expr_codegen`局限性
 
 1. `DAG`只能增加列无法删除。增加列时，遇到同名列会覆盖
@@ -169,7 +196,8 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 
 ## 特别语法
 
-1. 支持`C?T:F`三元表达式（仅可字符串中使用），底层会先转成`C or True if( T )else F`，然后修正成`T if C else F`，最后转成`if_else(C,T,F)`。支持与`if else`混用
+1. 支持`C?T:F`三元表达式（仅可字符串中使用），底层会先转成`C or True if( T )else F`，然后修正成`T if C else F`
+   ，最后转成`if_else(C,T,F)`。支持与`if else`混用
 2. `(A<B)*-1`,底层将转换成`int_(A<B)*-1`
 3. 为防止`A==B`被`sympy`替换成`False`，底层会换成`Eq(A,B)`
 4. `A^B`的含义与`convert_xor`参数有关，`convert_xor=True`底层会转换成`Pow(A,B)`，反之为`Xor(A,B)`。默认为`False`，用`**`表示乘方
@@ -179,6 +207,7 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 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()`元组解包，在底层会替换成
+11. 单行注释支持参数输入，如：`# --over_null`、`# --over_null=order_by`、`# --over_null=partition_by`
 
 ```python
 _x_0 = MACD()
@@ -191,7 +220,8 @@ C = unpack(_x_0, 2)
 
 1. 输出的数据，所有以`_`开头的列，最后会被自动删除。所以需要保留的变量一定不要以`_`开头
 2. 为减少重复计算，自动添加了了中间变量，以`_x_`开头，如`_x_0`，`_x_1`等。最后会被自动删除
-3. 单行表达式过长，或有重复计算，可以通过中间变量，将单行表达式改成多行。如果中间变量使用`_`开头，将会自动添加数字后缀，形成不同的变量，如`_A`会替换成`_A_0_`、`_A_1_`等。使用场景如下：
+3. 单行表达式过长，或有重复计算，可以通过中间变量，将单行表达式改成多行。如果中间变量使用`_`
+   开头，将会自动添加数字后缀，形成不同的变量，如`_A`会替换成`_A_0_`、`_A_1_`等。使用场景如下：
     1. 同一变量名，重复使用。本质是不同的变量
     2. 循环赋值，但`DAG`不支持有环。`=`号左右的同名变量其实是不同变量
 

expr_codegen-0.11.0/expr_codegen/_version.py

@@ -0,0 +1 @@
+__version__ = "0.11.0"

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/codes.py

@@ -2,10 +2,11 @@ import ast
 import re
 from ast import expr
 
+import ast_comments
 from black import Mode, format_str
 from sympy import Add, Mul, Pow, Eq, Not, Xor
 
-from expr_codegen.expr import register_symbols, dict_to_exprs
+from expr_codegen.expr import register_symbols, list_to_exprs
 
 
 class SyntaxTransformer(ast.NodeTransformer):
@@ -108,7 +109,8 @@ class SyntaxTransformer(ast.NodeTransformer):
     def visit_Subscript(self, node):
         if isinstance(node.slice, ast.Constant) and node.slice.value == 0:
             node = node.value
-        elif isinstance(node.slice, ast.UnaryOp) and isinstance(node.slice.operand, ast.Constant) and node.slice.operand.value == 0:
+        elif isinstance(node.slice, ast.UnaryOp) and isinstance(node.slice.operand,
+                                                                ast.Constant) and node.slice.operand.value == 0:
             node = node.value
         else:
             node = ast.Call(
@@ -328,6 +330,21 @@ def assigns_to_dict(assigns):
     return {ast.unparse(a.targets): ast.unparse(a.value) for a in assigns}
 
 
+def assigns_to_list(assigns):
+    """赋值表达式转成列表"""
+    outputs = []
+    for i, a in enumerate(assigns):
+        comment = "#"
+        if i + 1 < len(assigns):
+            b = assigns[i + 1]
+            if isinstance(b, ast_comments.Comment):
+                # comment = ast_comments.unparse(b)
+                comment = b.value
+        if isinstance(a, ast.Assign):
+            outputs.append((ast.unparse(a.targets), ast.unparse(a.value), comment))
+    return outputs
+
+
 def raw_to_code(raw):
     """导入语句转字符列表"""
     return '\n'.join([ast.unparse(a) for a in raw])
@@ -338,7 +355,7 @@ def sources_to_asts(*sources, convert_xor: bool):
 
     def _source_to_asts(source):
         """源代码"""
-        tree = ast.parse(source_replace(source))
+        tree = ast_comments.parse(source_replace(source))
 
         if isinstance(tree.body[0], ast.FunctionDef):
             body = tree.body[0].body
@@ -347,7 +364,7 @@ def sources_to_asts(*sources, convert_xor: bool):
 
         return body
 
-    tree = ast.parse("")
+    tree = ast_comments.parse("")
     for arg in sources:
         tree.body.extend(_source_to_asts(arg))
 
@@ -359,16 +376,21 @@ def sources_to_asts(*sources, convert_xor: bool):
     raw = []
     assigns = []
 
-    for node in tree.body:
+    for i, node in enumerate(tree.body):
         # 特殊处理的节点
         if isinstance(node, ast.Assign):
             assigns.append(node)
             continue
+        if isinstance(node, ast_comments.Comment):
+            # 添加注释
+            if node.inline and isinstance(assigns[-1], ast.Assign):
+                assigns.append(node)
+                continue
         # TODO 是否要把其它语句也加入？是否有安全问题？
         if isinstance(node, (ast.Import, ast.ImportFrom)):
             raw.append(node)
             continue
-    return raw_to_code(raw), assigns_to_dict(assigns), t.funcs_new, t.args_new, t.targets_new
+    return raw_to_code(raw), assigns_to_list(assigns), t.funcs_new, t.args_new, t.targets_new
 
 
 def _add_default_type(globals_):
@@ -394,5 +416,4 @@ def sources_to_exprs(globals_, *sources, convert_xor: bool):
     register_symbols(funcs_new, globals_, is_function=True)
     register_symbols(args_new, globals_, is_function=False)
     register_symbols(targets_new, globals_, is_function=False)
-    exprs_dict = dict_to_exprs(assigns, globals_)
-    return raw, exprs_dict
+    return raw, list_to_exprs(assigns, globals_)

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/expr.py

@@ -46,9 +46,8 @@ def register_symbols(syms, globals_, is_function: bool):
     return globals_
 
 
-def dict_to_exprs(exprs_src, globals_):
-    exprs_src = {k: sympify(v, globals_, evaluate=False) for k, v in exprs_src.items()}
-    return exprs_src
+def list_to_exprs(exprs_src, globals_):
+    return [(a, sympify(b, globals_, evaluate=False), c) for a, b, c in exprs_src]
 
 
 def append_node(node, output_exprs):
@@ -97,6 +96,12 @@ def get_symbols(expr, syms=None, return_str=True):
                 syms.append(arg.name)
             else:
                 syms.append(arg)
+        elif arg.is_Number:
+            # alpha_001 = log(1)+1
+            if return_str:
+                syms.append(str(arg))
+            else:
+                syms.append(arg)
         else:
             get_symbols(arg, syms, return_str)
     return syms
@@ -284,15 +289,15 @@ def get_key(children):
 def replace_exprs(exprs):
     """使用替换的方式简化表达式"""
     # Alpha101中大量ts_sum(x, 10)/10, 转成ts_mean(x, 10)
-    exprs = {k: _replace__ts_sum__to__ts_mean(v) for k, v in exprs.items()}
+    exprs = [(a, _replace__ts_sum__to__ts_mean(b), c) for a, b, c in exprs]
     # alpha_031中大量cs_rank(cs_rank(x)) 转成cs_rank(x)
-    exprs = {k: _replace__repeat(v) for k, v in exprs.items()}
+    exprs = [(a, _replace__repeat(b), c) for a, b, c in exprs]
     # 1.0*VWAP转VWAP
-    exprs = {k: _replace__one_mul(v) for k, v in exprs.items()}
+    exprs = [(a, _replace__one_mul(b), c) for a, b, c in exprs]
     # 将部分参数为1的ts函数进行简化
-    exprs = {k: _replace__ts_xxx_1(v) for k, v in exprs.items()}
+    exprs = [(a, _replace__ts_xxx_1(b), c) for a, b, c in exprs]
     # ts_delay转成ts_delta
-    exprs = {k: _replace__ts_delay__to__ts_delta(v) for k, v in exprs.items()}
+    exprs = [(a, _replace__ts_delay__to__ts_delta(b), c) for a, b, c in exprs]
 
     return exprs
 
@@ -435,34 +440,34 @@ def _replace__ts_delay__to__ts_delta(e):
         e = e.xreplace({node: replacement})
     return e
 
-
-def is_meaningless(e):
-    if _meaningless__ts_xxx_1(e):
-        return True
-    if _meaningless__xx_xx(e):
-        return True
-    return False
-
-
-def _meaningless__ts_xxx_1(e):
-    """ts_xxx部分函数如果参数为1，可直接丢弃"""
-    for node in preorder_traversal(e):
-        if len(node.args) >= 2:
-            if node.args[-1] == 1:
-                node_name = get_node_name(node)
-                if node_name in ('ts_delay', 'ts_delta', 'max_', 'min_'):
-                    return False
-                else:
-                    # 其它算子，参数1都认为无意义
-                    return True
-
-    return False
-
-
-def _meaningless__xx_xx(e):
-    """部分函数如果两参数完全一样，可直接丢弃"""
-    for node in preorder_traversal(e):
-        if len(node.args) >= 2:
-            if node.args[0] == node.args[1]:
-                return True
-    return False
+# def is_meaningless(e):
+#     if _meaningless__ts_xxx_1(e):
+#         return True
+#     if _meaningless__xx_xx(e):
+#         return True
+#     return False
+#
+#
+# def _meaningless__ts_xxx_1(e):
+#     """ts_xxx部分函数如果参数为1，可直接丢弃"""
+#     for node in preorder_traversal(e):
+#         if len(node.args) >= 2:
+#             node_name = get_node_name(node)
+#             if node_name in ('ts_delay', 'ts_delta'):
+#                 if not node.args[1].is_Integer:
+#                     return True
+#             if node_name.startswith('ts_'):
+#                 if not node.args[-1].is_Number:
+#                     return True
+#                 if node.args[-1] <= 1:
+#                     return True
+#     return False
+#
+#
+# def _meaningless__xx_xx(e):
+#     """部分函数如果两参数完全一样，可直接丢弃"""
+#     for node in preorder_traversal(e):
+#         if len(node.args) >= 2:
+#             if node.args[0] == node.args[1]:
+#                 return True
+#     return False

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/model.py

@@ -4,7 +4,7 @@ from itertools import product
 import networkx as nx
 from sympy import symbols
 
-from expr_codegen.dag import zero_indegree, hierarchy_pos, remove_paths_by_zero_outdegree
+from expr_codegen.dag import zero_indegree, hierarchy_pos, remove_paths_by_zero_outdegree, zero_outdegree
 from expr_codegen.expr import CL, get_symbols, get_children, get_key, is_simple_expr
 
 _RESERVED_WORD_ = {'_NONE_', '_TRUE_', '_FALSE_'}
@@ -196,15 +196,15 @@ def create_dag_exprs(exprs):
     # 创建有向无环图
     G = nx.DiGraph()
 
-    for symbol, expr in exprs.items():
+    for symbol, expr, comment in exprs:
         # if symbol.name == 'GP_0':
         #     test = 1
         if expr.is_Symbol:
-            G.add_node(symbol.name, symbol=symbol, expr=expr)
+            G.add_node(symbol.name, symbol=symbol, expr=expr, comment=comment)
             G.add_edge(expr.name, symbol.name)
         else:
             # 添加中间节点
-            G.add_node(symbol.name, symbol=symbol, expr=expr)
+            G.add_node(symbol.name, symbol=symbol, expr=expr, comment=comment)
             syms = get_symbols(expr, return_str=True)
             for sym in syms:
                 # 由于边的原因，这里会主动生成一些源节点
@@ -221,6 +221,10 @@ def create_dag_exprs(exprs):
         s = symbols(node)
         G.nodes[node]['symbol'] = s
         G.nodes[node]['expr'] = s
+        G.nodes[node]['comment'] = "#"
+    #
+    # for node in zero_outdegree(G):
+    #     print(11, G.nodes[node]['comment'])
     return G
 
 
@@ -380,9 +384,9 @@ def skip_expr_node(G: nx.DiGraph, node, keep_nodes):
     return G
 
 
-def dag_start(exprs_dict, func, func_kwargs, date, asset):
+def dag_start(exprs_list, func, func_kwargs, date, asset):
     """初始生成DAG"""
-    G = create_dag_exprs(exprs_dict)
+    G = create_dag_exprs(exprs_list)
     G = init_dag_exprs(G, func, func_kwargs, date, asset)
 
     # 分层输出
@@ -413,11 +417,12 @@ def dag_end(G):
         for node in generation:
             key = G.nodes[node]['key']
             expr = G.nodes[node]['expr']
+            comment = G.nodes[node]['comment']
             symbols = G.nodes[node]['symbols']
             # 这几个特殊的不算成字段名
             symbols = list(set(symbols) - _RESERVED_WORD_)
 
-            exprs_ldl.append(key, (node, expr, symbols))
+            exprs_ldl.append(key, (node, expr, symbols, comment))
 
     exprs_ldl._list = exprs_ldl.values()[1:]
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/pandas/code.py

@@ -67,9 +67,9 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                     func_code.append(f"    # " + '=' * 40)
                     exprs_dst.append(f"#" + '=' * 40 + func_name)
                 else:
-                    va, ex, sym = kv
+                    va, ex, sym, comment = kv
                     func_code.append(f"    # {va} = {ex}\n    g[{va}] = {p.doprint(ex)}")
-                    exprs_dst.append(f"{va} = {ex}")
+                    exprs_dst.append(f"{va} = {ex} {comment}")
                     if va not in syms_dst:
                         syms_out.append(va)
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/pandas/template.py.j2

@@ -42,8 +42,8 @@ def {{ key }}(df: pd.DataFrame) -> pd.DataFrame:
 """
 
 """
-{%-for key, value in exprs_src.items() %}
-{{ key }} = {{ value-}}
+{%-for a,b,c in exprs_src %}
+{{ a }} = {{ b}} {{c-}}
 {% endfor %}
 """
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/polars_group/code.py

@@ -70,7 +70,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                     func_code.append(f"    df = df.with_columns(")
                     exprs_dst.append(f"#" + '=' * 40 + func_name)
                 else:
-                    va, ex, sym = kv
+                    va, ex, sym, comment = kv
                     s1 = str(ex)
                     s2 = p.doprint(ex)
                     if s1 != s2:
@@ -78,7 +78,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                         func_code.append(f"# {va} = {s1}")
 
                     func_code.append(f"{va}={s2},")
-                    exprs_dst.append(f"{va} = {s1}")
+                    exprs_dst.append(f"{va} = {s1} {comment}")
                     if va not in syms_dst:
                         syms_out.append(va)
             func_code.append(f"    )")

{expr_codegen-0.10.15/expr_codegen/polars_over → expr_codegen-0.11.0/expr_codegen/polars_group}/template.py.j2

@@ -51,8 +51,8 @@ def {{ key }}(df: DataFrame) -> DataFrame:
 """
 
 """
-{%-for key, value in exprs_src.items() %}
-{{ key }} = {{ value-}}
+{%-for a,b,c in exprs_src %}
+{{ a }} = {{ b}} {{c-}}
 {% endfor %}
 """
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/polars_over/code.py

@@ -1,3 +1,4 @@
+import argparse
 import os
 from typing import Sequence, Dict, Literal
 
@@ -43,6 +44,9 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
             over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
             **kwargs):
     """基于模板的代码生成"""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--over_null", type=str, nargs="?", default=over_null)
+
     # 打印Polars风格代码
     p = PolarsStrPrinter()
 
@@ -71,7 +75,9 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                     func_code.append(f"    df = df.with_columns(")
                     exprs_dst.append(f"#" + '=' * 40 + func_name)
                 else:
-                    va, ex, sym = kv
+                    va, ex, sym, comment = kv
+                    # 多个#时，只取第一个#后的参数
+                    args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
                     s1 = str(ex)
                     s2 = p.doprint(ex)
                     if s1 != s2:
@@ -84,9 +90,9 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                             _sym = f"pl.all_horizontal({','.join(_sym)})"
                         else:
                             _sym = ','.join(_sym)
-                        if over_null == 'partition_by':
+                        if args.over_null == 'partition_by':
                             func_code.append(f"{va}=({s2}).over({_sym}, _ASSET_, order_by=_DATE_),")
-                        elif over_null == 'order_by':
+                        elif args.over_null == 'order_by':
                             func_code.append(f"{va}=({s2}).over(_ASSET_, order_by=[{_sym}, _DATE_]),")
                         else:
                             func_code.append(f"{va}=({s2}).over(_ASSET_, order_by=_DATE_),")
@@ -96,7 +102,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                         func_code.append(f"{va}=({s2}).over(_DATE_, '{k[2]}'),")
                     else:
                         func_code.append(f"{va}={s2},")
-                    exprs_dst.append(f"{va} = {s1}")
+                    exprs_dst.append(f"{va} = {s1} {comment}")
                     if va not in syms_dst:
                         syms_out.append(va)
             func_code.append(f"    )")

{expr_codegen-0.10.15/expr_codegen/polars_group → expr_codegen-0.11.0/expr_codegen/polars_over}/template.py.j2

@@ -51,8 +51,8 @@ def {{ key }}(df: DataFrame) -> DataFrame:
 """
 
 """
-{%-for key, value in exprs_src.items() %}
-{{ key }} = {{ value-}}
+{%-for a,b,c in exprs_src %}
+{{ a }} = {{ b}} {{c-}}
 {% endfor %}
 """
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen/tool.py

@@ -61,7 +61,7 @@ class ExprTool:
     def __init__(self):
         self.get_current_func = get_current_by_prefix
         self.get_current_func_kwargs = {}
-        self.exprs_dict = {}
+        self.exprs_list = {}
         self.exprs_names = []
         self.globals_ = {}
 
@@ -92,7 +92,7 @@ class ExprTool:
         # print(exprs)
         return exprs, syms
 
-    def merge(self, date, asset, **kwargs):
+    def merge(self, date, asset, args):
         """合并多个表达式
 
         1. 先抽取分割子公式
@@ -100,28 +100,31 @@ class ExprTool:
 
         Parameters
         ----------
-        kwargs
-            表达式字典
+        args
+            表达式列表
 
         Returns
         -------
         表达式列表
         """
         # 抽取前先化简
-        kwargs = {k: simplify2(v) for k, v in kwargs.items()}
+        args = [(a, simplify2(b), c) for a, b, c in args]
 
-        exprs_syms = [self.extract(v, date, asset) for v in kwargs.values()]
+        # 保留了注释信息
+        exprs_syms = [(self.extract(b, date, asset), c) for a, b, c in args]
         exprs = []
         syms = []
-        for e, s in exprs_syms:
-            exprs.extend(e)
+        for (e, s), c in exprs_syms:
             syms.extend(s)
+            for _ in e:
+                # 抽取的表达式添加注释
+                exprs.append((_, c))
 
         syms = sorted(set(syms), key=syms.index)
         # 如果目标有重复表达式，这里会混乱
         exprs = sorted(set(exprs), key=exprs.index)
         # 这里不能合并简化与未简化的表达式，会导致cse时失败，需要简化表达式合并
-        exprs = exprs + list(kwargs.values())
+        exprs = exprs + [(b, c) for a, b, c in args]
 
         # print(exprs)
         syms = [str(s) for s in syms]
@@ -130,18 +133,18 @@ class ExprTool:
     def reduce(self, repl, redu):
         """减少中间变量数量，有利用减少内存占用"""
 
-        exprs_dict = {}
+        exprs_list = []
 
         # cse前简化一次，cse后不再简化
         # (~开盘涨停 & 昨收涨停) | (~收盘涨停 & 最高涨停)
-        for variable, expr in repl:
-            exprs_dict[variable] = expr
-        for variable, expr in redu:
-            exprs_dict[variable] = expr
+        for a, b in repl:
+            exprs_list.append((a, b, "#"))
+        for a, b, c in redu:
+            exprs_list.append((a, b, c))
 
-        return exprs_dict
+        return exprs_list
 
-    def cse(self, exprs, symbols_repl=None, symbols_redu=None):
+    def cse(self, exprs, symbols_repl=None, exprs_src=None):
         """多个子公式+长公式，提取公共公式
 
         Parameters
@@ -150,7 +153,7 @@ class ExprTool:
             表达式列表
         symbols_repl
             中间字段名迭代器
-        symbols_redu
+        exprs_src
             最终字段名列表
 
         Returns
@@ -163,34 +166,38 @@ class ExprTool:
             表达式
 
         """
-        self.exprs_names = list(symbols_redu)
+        self.exprs_names = [a for a, b, c in exprs_src]
+        # 包含了注释信息
+        _exprs = [a for a, b in exprs]
 
-        repl, redu = cse(exprs, symbols_repl, optimizations="basic")
-        outputs_len = len(symbols_redu)
+        # 注意：对于表达式右边相同，左边不同的情况，会当成一个处理
+        repl, redu = cse(_exprs, symbols_repl, optimizations="basic")
+        outputs_len = len(exprs_src)
 
         new_redu = []
-        symbols_redu = iter(symbols_redu)
+        symbols_redu = iter(exprs_src)
         for expr in redu[-outputs_len:]:
             # 可能部分表达式只在之前出现过，后面完全用不到如，ts_rank(ts_decay_linear(x_147, 11.4157), 6.72611)
             variable = next(symbols_redu)
-            variable = symbols(variable)
-            new_redu.append((variable, expr))
+            a = symbols(variable[0])
+            new_redu.append((a, expr, variable[2]))
 
-        self.exprs_dict = self.reduce(repl, new_redu)
+        self.exprs_list = self.reduce(repl, new_redu)
 
         # with open("exprs.pickle", "wb") as file:
         #     pickle.dump(exprs_dict, file)
 
-        return self.exprs_dict
+        return self.exprs_list
 
     def dag(self, merge: bool, date, asset):
         """生成DAG"""
-        G = dag_start(self.exprs_dict, self.get_current_func, self.get_current_func_kwargs, date, asset)
+        G = dag_start(self.exprs_list, self.get_current_func, self.get_current_func_kwargs, date, asset)
         if merge:
             G = dag_middle(G, self.exprs_names, self.get_current_func, self.get_current_func_kwargs, date, asset)
         return dag_end(G)
 
-    def all(self, exprs_src, style: Literal['pandas', 'polars_group', 'polars_over'] = 'polars_over', template_file: str = 'template.py.j2',
+    def all(self, exprs_src, style: Literal['pandas', 'polars_group', 'polars_over'] = 'polars_over',
+            template_file: str = 'template.py.j2',
             replace: bool = True, regroup: bool = False, format: bool = True,
             date='date', asset='asset',
             alias: Dict[str, str] = {},
@@ -200,8 +207,8 @@ class ExprTool:
 
         Parameters
         ----------
-        exprs_src: dict
-            表达式字典
+        exprs_src: list
+            表达式列表
         style: str
             代码风格。可选值 ('polars_group', 'polars_over', 'pandas')
         template_file: str
@@ -232,11 +239,11 @@ class ExprTool:
             exprs_src = replace_exprs(exprs_src)
 
         # 子表达式在前，原表式在最后
-        exprs_dst, syms_dst = self.merge(date, asset, **exprs_src)
+        exprs_dst, syms_dst = self.merge(date, asset, exprs_src)
         syms_dst = list(set(syms_dst) - _RESERVED_WORD_)
 
         # 提取公共表达式
-        self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), symbols_redu=exprs_src.keys())
+        self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), exprs_src=exprs_src)
         # 有向无环图流转
         exprs_ldl, G = self.dag(True, date, asset)
 
@@ -272,14 +279,15 @@ class ExprTool:
                   extra_codes: str,
                   output_file: str,
                   convert_xor: bool,
-                  style: Literal['pandas', 'polars_group', 'polars_over'] = 'polars_over', template_file: str = 'template.py.j2',
+                  style: Literal['pandas', 'polars_group', 'polars_over'] = 'polars_over',
+                  template_file: str = 'template.py.j2',
                   date: str = 'date', asset: str = 'asset',
                   **kwargs) -> str:
         """通过字符串生成代码， 加了缓存，多次调用不重复生成"""
-        raw, exprs_dict = sources_to_exprs(self.globals_, source, *more_sources, convert_xor=convert_xor)
+        raw, exprs_list = sources_to_exprs(self.globals_, source, *more_sources, convert_xor=convert_xor)
 
         # 生成代码
-        code, G = _TOOL_.all(exprs_dict, style=style, template_file=template_file,
+        code, G = _TOOL_.all(exprs_list, style=style, template_file=template_file,
                              replace=True, regroup=True, format=True,
                              date=date, asset=asset,
                              # 复制了需要使用的函数，还复制了最原始的表达式
@@ -333,6 +341,7 @@ _TOOL_ = ExprTool()
 
 def codegen_exec(df: Optional[DataFrame],
                  *codes,
+                 over_null: Literal['partition_by', 'order_by', None],
                  extra_codes: str = r'CS_SW_L1 = r"^sw_l1_\d+$"',
                  output_file: Union[str, TextIOBase, None] = None,
                  run_file: Union[bool, str] = False,
@@ -340,7 +349,7 @@ def codegen_exec(df: Optional[DataFrame],
                  style: Literal['pandas', 'polars_group', 'polars_over'] = 'polars_over',
                  template_file: str = 'template.py.j2',
                  date: str = 'date', asset: str = 'asset',
-                 over_null: Literal['partition_by', 'order_by', None] = 'partition_by',
+
                  **kwargs) -> Optional[DataFrame]:
     """快速转换源代码并执行
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: expr_codegen
-Version: 0.10.15
+Version: 0.11.0
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License
@@ -43,6 +43,7 @@ Requires-Dist: Jinja2
 Requires-Dist: networkx
 Requires-Dist: loguru
 Requires-Dist: sympy
+Requires-Dist: ast-comments
 Provides-Extra: streamlit
 Requires-Dist: streamlit; extra == "streamlit"
 Requires-Dist: streamlit-ace; extra == "streamlit"
@@ -211,6 +212,33 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 2. `over_null='order_by'`。分到一个区域，`null`排在前面
 3. `over_null=None`。不处理，直接调用，速度更快。如果确信不会中段产生`null`建议使用此参数
 
+`codegen_exec(over_null='partition_by')`为全局使用`partition_by`。但遇到`ts_count_nulls`这类`null`
+函数就得使用`over_null=None`，所以本工具还新添了注释功能来指定单行表达式参数
+
+1. `# --over_null partition_by`。单行`over_null='partition_by'`
+2. `# --over_null=order_by`。单行`over_null='order_by'`
+3. `# --over_null`。单行`over_null=None`
+4. `# `。取`codegen_exec`参数传入的`over_null`值
+
+注意：
+
+1. `# --over_null`传参注释只能写在单行表达式的后面，不能独立成一行，否则会被忽略
+2. `# --over_null # --over_null=order_by`多个`#`时，只取第一个有效
+3. 只对最外层`ts`函数有效。如果`ts`函数不在外层，需要人工提炼。例如：
+   ```python
+   X1 = cs_rank(ts_mean(CLOSE, 3)) # --over_null=order_by # 应用在cs_rank上，没有意义
+   X2 = ts_rank(ts_mean(CLOSE, 3), 5) # --over_null=order_by # 本以为应用在ts_rank(ts_mean)上，但由于出现了公共ts_mean，其实是应用在ts_rank(_x_0)上
+   ```
+
+   需写成
+
+   ```python
+   _x_0 = ts_mean(CLOSE, 3)  # --over_null=order_by 
+   X1 = cs_rank(_x_0)
+   X2 = ts_rank(_x_0, 5)
+   ```
+4. 由于很容易搞错，强烈建议生成`output_file`,检查生成的代码是否正确。
+
 ## `expr_codegen`局限性
 
 1. `DAG`只能增加列无法删除。增加列时，遇到同名列会覆盖
@@ -220,7 +248,8 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 
 ## 特别语法
 
-1. 支持`C?T:F`三元表达式（仅可字符串中使用），底层会先转成`C or True if( T )else F`，然后修正成`T if C else F`，最后转成`if_else(C,T,F)`。支持与`if else`混用
+1. 支持`C?T:F`三元表达式（仅可字符串中使用），底层会先转成`C or True if( T )else F`，然后修正成`T if C else F`
+   ，最后转成`if_else(C,T,F)`。支持与`if else`混用
 2. `(A<B)*-1`,底层将转换成`int_(A<B)*-1`
 3. 为防止`A==B`被`sympy`替换成`False`，底层会换成`Eq(A,B)`
 4. `A^B`的含义与`convert_xor`参数有关，`convert_xor=True`底层会转换成`Pow(A,B)`，反之为`Xor(A,B)`。默认为`False`，用`**`表示乘方
@@ -230,6 +259,7 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
 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()`元组解包，在底层会替换成
+11. 单行注释支持参数输入，如：`# --over_null`、`# --over_null=order_by`、`# --over_null=partition_by`
 
 ```python
 _x_0 = MACD()
@@ -242,7 +272,8 @@ C = unpack(_x_0, 2)
 
 1. 输出的数据，所有以`_`开头的列，最后会被自动删除。所以需要保留的变量一定不要以`_`开头
 2. 为减少重复计算，自动添加了了中间变量，以`_x_`开头，如`_x_0`，`_x_1`等。最后会被自动删除
-3. 单行表达式过长，或有重复计算，可以通过中间变量，将单行表达式改成多行。如果中间变量使用`_`开头，将会自动添加数字后缀，形成不同的变量，如`_A`会替换成`_A_0_`、`_A_1_`等。使用场景如下：
+3. 单行表达式过长，或有重复计算，可以通过中间变量，将单行表达式改成多行。如果中间变量使用`_`
+   开头，将会自动添加数字后缀，形成不同的变量，如`_A`会替换成`_A_0_`、`_A_1_`等。使用场景如下：
     1. 同一变量名，重复使用。本质是不同的变量
     2. 循环赋值，但`DAG`不支持有环。`=`号左右的同名变量其实是不同变量
 

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/expr_codegen.egg-info/requires.txt

@@ -3,6 +3,7 @@ Jinja2
 networkx
 loguru
 sympy
+ast-comments
 
 [streamlit]
 streamlit

{expr_codegen-0.10.15 → expr_codegen-0.11.0}/pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     'networkx',
     'loguru',
     'sympy',
+    'ast-comments',
     # 'polars_ta',
 ]
 dynamic = ["version"]

expr_codegen-0.10.15/expr_codegen/_version.py

@@ -1 +0,0 @@
-__version__ = "0.10.15"