expr-codegen 0.16.6__tar.gz → 0.17.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (33) hide show
  1. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/PKG-INFO +28 -18
  2. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/README.md +26 -16
  3. expr_codegen-0.17.0/expr_codegen/_version.py +1 -0
  4. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/expr.py +11 -7
  5. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/model.py +30 -4
  6. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/code.py +3 -3
  7. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/code.py +4 -4
  8. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/code.py +1 -1
  9. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/tool.py +44 -50
  10. expr_codegen-0.16.6/expr_codegen/_version.py +0 -1
  11. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/.gitignore +0 -0
  12. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/LICENSE +0 -0
  13. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/__init__.py +0 -0
  14. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/codes.py +0 -0
  15. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/dag.py +0 -0
  16. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/latex/__init__.py +0 -0
  17. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/latex/printer.py +0 -0
  18. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/__init__.py +0 -0
  19. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/code.py +0 -0
  20. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/helper.py +0 -0
  21. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/printer.py +0 -0
  22. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/ta.py +0 -0
  23. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/pandas/template.py.j2 +0 -0
  24. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/__init__.py +0 -0
  25. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/printer.py +0 -0
  26. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/polars/template.py.j2 +0 -0
  27. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/__init__.py +0 -0
  28. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/printer.py +0 -0
  29. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/rust/template.rs.j2 +0 -0
  30. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/__init__.py +0 -0
  31. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/printer.py +0 -0
  32. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/expr_codegen/sql/template.sql.j2 +0 -0
  33. {expr_codegen-0.16.6 → expr_codegen-0.17.0}/pyproject.toml +0 -0
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: expr_codegen
3
- Version: 0.16.6
3
+ Version: 0.17.0
4
4
  Summary: symbol expression to polars expression tool
5
5
  Author-email: wukan <wu-kan@163.com>
6
6
  License: BSD 3-Clause License
@@ -116,19 +116,19 @@ def _code_block_2():
116
116
  CPV = cs_zscore(_corr) + cs_zscore(_beta)
117
117
 
118
118
 
119
- code = codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by', output_file=sys.stdout) # 打印代码
120
- code = codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by', output_file="output.py") # 保存到文件
121
- code = codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by') # 只执行,不保存代码
119
+ code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=sys.stdout) # 打印代码
120
+ code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file="output.py") # 保存到文件
121
+ code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by') # 只执行,不保存代码
122
122
 
123
123
  code = StringIO()
124
- codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by', output_file=code) # 保存到字符串
124
+ codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=code) # 保存到字符串
125
125
  code.seek(0)
126
126
  code.read() # 读取代码
127
127
 
128
128
  # TODO 替换成合适的数据
129
129
  df = pl.DataFrame()
130
- df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='partition_by').collect() # Lazy CPU
131
- df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='partition_by').collect(engine="gpu") # Lazy GPU
130
+ df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect() # Lazy CPU
131
+ df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect(engine="gpu") # Lazy GPU
132
132
 
133
133
  ```
134
134
 
@@ -217,11 +217,11 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
217
217
 
218
218
  第2种开头的`null`区域,是否影响结果由算子所决定,特别时是多列输入时`null`区域可能有数据
219
219
 
220
- 1. `over_null='partition_by'`。分到两个区域
221
- 2. `over_null='order_by'`。分到一个区域,`null`排在前面
220
+ 1. `over_null='partition_by'`。分到两个区域。多`ts`嵌套时,分区导致数据区变短,输出的有效数据变得更短
221
+ 2. `over_null='order_by'`。分到一个区域,`null`排在前面。多`ts`嵌套时,数据区长度不变,输出的有效数据符合算法
222
222
  3. `over_null=None`。不处理,直接调用,速度更快。如果确信不会中段产生`null`建议使用此参数
223
223
 
224
- `codegen_exec(over_null='partition_by')`为全局使用`partition_by`。但遇到`ts_count_nulls`这类`null`
224
+ `codegen_exec(over_null='order_by')`为全局使用`order_by`。但遇到`ts_count_nulls`这类`null`
225
225
  函数就得使用`over_null=None`,所以本工具还新添了注释功能来指定单行表达式参数
226
226
 
227
227
  1. `# --over_null partition_by`。单行`over_null='partition_by'`
@@ -233,18 +233,28 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
233
233
 
234
234
  1. `# --over_null`传参注释只能写在单行表达式的后面,不能独立成一行,否则会被忽略
235
235
  2. `# --over_null # --over_null=order_by`多个`#`时,只取第一个有效
236
- 3. 只对最外层`ts`函数有效。如果`ts`函数不在外层,需要人工提炼。例如:
236
+ 3. 只对`ts`函数有效。并能向上游传播。例如:
237
237
  ```python
238
- X1 = cs_rank(ts_mean(CLOSE, 3)) # --over_null=order_by # 应用在cs_rank上,没有意义
239
- X2 = ts_rank(ts_mean(CLOSE, 3), 5) # --over_null=order_by # 本以为应用在ts_rank(ts_mean)上,但由于出现了公共ts_mean,其实是应用在ts_rank(_x_0)上
238
+ X1 = ts_mean(ts_mean(CLOSE, 2), 2) # --over_null order_by
239
+ X2 = ts_mean(CLOSE, 2) # --over_null=order_by
240
+ X3 = ts_mean(CLOSE, 2) # 取`codegen_exec`参数传入的`over_null`值
241
+ X4 = cs_rank(CLOSE) + ts_mean(CLOSE, 2) # --over_null None
240
242
  ```
241
243
 
242
- 需写成
244
+ 会自动翻译成
243
245
 
244
246
  ```python
245
- _x_0 = ts_mean(CLOSE, 3) # --over_null=order_by
246
- X1 = cs_rank(_x_0)
247
- X2 = ts_rank(_x_0, 5)
247
+ #========================================func_0_cs__date
248
+ _x_1 = cs_rank(CLOSE) # --over_null=None
249
+ #========================================func_0_ts__asset
250
+ _x_0 = ts_mean(CLOSE, 2) # --over_null=order_by
251
+ _x_2 = ts_mean(CLOSE, 2) # --over_null=None
252
+ #========================================func_0_ts__asset
253
+ X1 = ts_mean(_x_0, 2) # --over_null=order_by
254
+ #========================================func_1_cl
255
+ X2 = _x_0 # --over_null=order_by
256
+ X3 = _x_0 # --over_null=order_by
257
+ X4 = _x_1 + _x_2 # --over_null=None
248
258
  ```
249
259
  4. 由于很容易搞错,强烈建议生成`output_file`,检查生成的代码是否正确。
250
260
 
@@ -66,19 +66,19 @@ def _code_block_2():
66
66
  CPV = cs_zscore(_corr) + cs_zscore(_beta)
67
67
 
68
68
 
69
- code = codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by', output_file=sys.stdout) # 打印代码
70
- code = codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by', output_file="output.py") # 保存到文件
71
- code = codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by') # 只执行,不保存代码
69
+ code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=sys.stdout) # 打印代码
70
+ code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file="output.py") # 保存到文件
71
+ code = codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by') # 只执行,不保存代码
72
72
 
73
73
  code = StringIO()
74
- codegen_exec(None, _code_block_1, _code_block_2, over_null='partition_by', output_file=code) # 保存到字符串
74
+ codegen_exec(None, _code_block_1, _code_block_2, over_null='order_by', output_file=code) # 保存到字符串
75
75
  code.seek(0)
76
76
  code.read() # 读取代码
77
77
 
78
78
  # TODO 替换成合适的数据
79
79
  df = pl.DataFrame()
80
- df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='partition_by').collect() # Lazy CPU
81
- df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='partition_by').collect(engine="gpu") # Lazy GPU
80
+ df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect() # Lazy CPU
81
+ df = codegen_exec(df.lazy(), _code_block_1, _code_block_2, over_null='order_by').collect(engine="gpu") # Lazy GPU
82
82
 
83
83
  ```
84
84
 
@@ -167,11 +167,11 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
167
167
 
168
168
  第2种开头的`null`区域,是否影响结果由算子所决定,特别时是多列输入时`null`区域可能有数据
169
169
 
170
- 1. `over_null='partition_by'`。分到两个区域
171
- 2. `over_null='order_by'`。分到一个区域,`null`排在前面
170
+ 1. `over_null='partition_by'`。分到两个区域。多`ts`嵌套时,分区导致数据区变短,输出的有效数据变得更短
171
+ 2. `over_null='order_by'`。分到一个区域,`null`排在前面。多`ts`嵌套时,数据区长度不变,输出的有效数据符合算法
172
172
  3. `over_null=None`。不处理,直接调用,速度更快。如果确信不会中段产生`null`建议使用此参数
173
173
 
174
- `codegen_exec(over_null='partition_by')`为全局使用`partition_by`。但遇到`ts_count_nulls`这类`null`
174
+ `codegen_exec(over_null='order_by')`为全局使用`order_by`。但遇到`ts_count_nulls`这类`null`
175
175
  函数就得使用`over_null=None`,所以本工具还新添了注释功能来指定单行表达式参数
176
176
 
177
177
  1. `# --over_null partition_by`。单行`over_null='partition_by'`
@@ -183,18 +183,28 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
183
183
 
184
184
  1. `# --over_null`传参注释只能写在单行表达式的后面,不能独立成一行,否则会被忽略
185
185
  2. `# --over_null # --over_null=order_by`多个`#`时,只取第一个有效
186
- 3. 只对最外层`ts`函数有效。如果`ts`函数不在外层,需要人工提炼。例如:
186
+ 3. 只对`ts`函数有效。并能向上游传播。例如:
187
187
  ```python
188
- X1 = cs_rank(ts_mean(CLOSE, 3)) # --over_null=order_by # 应用在cs_rank上,没有意义
189
- X2 = ts_rank(ts_mean(CLOSE, 3), 5) # --over_null=order_by # 本以为应用在ts_rank(ts_mean)上,但由于出现了公共ts_mean,其实是应用在ts_rank(_x_0)上
188
+ X1 = ts_mean(ts_mean(CLOSE, 2), 2) # --over_null order_by
189
+ X2 = ts_mean(CLOSE, 2) # --over_null=order_by
190
+ X3 = ts_mean(CLOSE, 2) # 取`codegen_exec`参数传入的`over_null`值
191
+ X4 = cs_rank(CLOSE) + ts_mean(CLOSE, 2) # --over_null None
190
192
  ```
191
193
 
192
- 需写成
194
+ 会自动翻译成
193
195
 
194
196
  ```python
195
- _x_0 = ts_mean(CLOSE, 3) # --over_null=order_by
196
- X1 = cs_rank(_x_0)
197
- X2 = ts_rank(_x_0, 5)
197
+ #========================================func_0_cs__date
198
+ _x_1 = cs_rank(CLOSE) # --over_null=None
199
+ #========================================func_0_ts__asset
200
+ _x_0 = ts_mean(CLOSE, 2) # --over_null=order_by
201
+ _x_2 = ts_mean(CLOSE, 2) # --over_null=None
202
+ #========================================func_0_ts__asset
203
+ X1 = ts_mean(_x_0, 2) # --over_null=order_by
204
+ #========================================func_1_cl
205
+ X2 = _x_0 # --over_null=order_by
206
+ X3 = _x_0 # --over_null=order_by
207
+ X4 = _x_1 + _x_2 # --over_null=None
198
208
  ```
199
209
  4. 由于很容易搞错,强烈建议生成`output_file`,检查生成的代码是否正确。
200
210
 
@@ -0,0 +1 @@
1
+ __version__ = "0.17.0"
@@ -200,7 +200,7 @@ def get_children(func, func_kwargs, expr, output_exprs, output_symbols, date, as
200
200
  expr
201
201
  表达式
202
202
  output_exprs
203
- 输出分割后的了表达式
203
+ 输出分割后的子表达式
204
204
  output_symbols
205
205
  输出每个子表达式中的符号
206
206
  date
@@ -236,17 +236,21 @@ def get_children(func, func_kwargs, expr, output_exprs, output_symbols, date, as
236
236
  # 外部与内部不同,需处理
237
237
  for i, child in enumerate(children):
238
238
  # 不在子中即表示不同
239
- if curr in child:
240
- continue
239
+ # if curr in child:
240
+ # continue
241
241
  if expr.args[i].is_Atom:
242
242
  # print(expr.args[i], 'is_Atom 2')
243
243
  continue
244
244
  output_exprs = append_node(expr.args[i], output_exprs)
245
245
  else:
246
- # ts_sum(OPEN, 5)*ts_sum(RETURNS, 5) ('cl',) [{('ts', 'asset', 'date')}, {('ts', 'asset', 'date')}] 6 alpha_008
247
- pass
248
- # if isinstance(expr, Mul):
249
- # output_exprs = append_node(expr, output_exprs)
246
+ for i, child in enumerate(children):
247
+ if not child: # 跳过空集合
248
+ continue
249
+ if next(iter(child))[0] == CL:
250
+ continue
251
+ if expr.args[i].is_Atom:
252
+ continue
253
+ output_exprs = append_node(expr.args[i], output_exprs)
250
254
 
251
255
  # 按需返回,当前是基础算子就返回下一层信息,否则返回当前
252
256
  if curr[0] == CL:
@@ -349,6 +349,8 @@ def skip_expr_node(G: nx.DiGraph, node, keep_nodes):
349
349
  e = G.nodes[s]['expr']
350
350
  e = e.xreplace({symbol: expr})
351
351
  G.nodes[s]['expr'] = e
352
+ if d['comment'] != "#":
353
+ G.nodes[s]['comment'] = d['comment']
352
354
 
353
355
  # 这里用了product生成多个关联边
354
356
  G.add_edges_from(product(pred, succ))
@@ -365,14 +367,11 @@ def dag_start(exprs_list, func, func_kwargs, date, asset):
365
367
  return G
366
368
 
367
369
 
368
- def dag_middle(G, exprs_names, skip_columns, func, func_kwargs, date, asset, skip_simplify):
370
+ def dag_middle(G, exprs_names, skip_columns, func, func_kwargs, date, asset):
369
371
  """删除几个没有必要的节点"""
370
372
  # 以下划线开头的节点,不保留
371
373
  keep_nodes = [k for k in exprs_names if not k.startswith('_')]
372
374
 
373
- if not skip_simplify:
374
- # ts_rank(-RET - -RET, 20),防止替换成ts_rank(0, 20)
375
- G = merge_nodes_1(G, keep_nodes, *keep_nodes)
376
375
  G = merge_nodes_2(G, keep_nodes, *keep_nodes)
377
376
 
378
377
  # 移除0出度的节点,但保留部分
@@ -405,3 +404,30 @@ def dag_end(G):
405
404
  exprs_ldl._list = exprs_ldl.values()[1:]
406
405
 
407
406
  return exprs_ldl, G
407
+
408
+
409
+ def propagate_comment_recursive(G):
410
+ """按拓扑排序反序遍历,确保从终点开始"""
411
+ from collections import deque
412
+
413
+ # Kahn算法求拓扑排序
414
+ in_degree = {node: G.in_degree(node) for node in G.nodes()}
415
+ queue = deque([node for node in G.nodes() if in_degree[node] == 0])
416
+ topo_order = []
417
+
418
+ while queue:
419
+ node = queue.popleft()
420
+ topo_order.append(node)
421
+ for succ in G.successors(node):
422
+ in_degree[succ] -= 1
423
+ if in_degree[succ] == 0:
424
+ queue.append(succ)
425
+
426
+ # 反序遍历拓扑排序(从终点开始)
427
+ for node in reversed(topo_order):
428
+ comment = G.nodes[node]['comment']
429
+ if comment == "#":
430
+ continue
431
+ for pred in G.predecessors(node):
432
+ if G.nodes[pred]['comment'] == "#":
433
+ G.nodes[pred]['comment'] = comment
@@ -44,7 +44,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
44
44
  filename,
45
45
  date='date', asset='asset',
46
46
  extra_codes: Sequence[str] = (),
47
- over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
47
+ over_null: Literal['order_by', 'partition_by', None] = 'order_by',
48
48
  ge_date_idx: int = 0,
49
49
  **kwargs):
50
50
  """基于模板的代码生成"""
@@ -83,14 +83,14 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
83
83
  exprs_dst.append(f"#" + '=' * 40 + func_name)
84
84
  else:
85
85
  va, ex, sym, comment = kv
86
- # 多个#时,只取第一个#后的参数
87
- args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
88
86
  s1 = str(ex)
89
87
  s2 = p.doprint(ex)
90
88
  if s1 != s2:
91
89
  # 不想等,打印注释,显示会更直观察
92
90
  func_code.append(f"# {va} = {s1}")
93
91
  if k[0] == TS:
92
+ # 多个#时,只取第一个#后的参数
93
+ args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
94
94
  ts_func_name = func_name
95
95
  # https://github.com/pola-rs/polars/issues/12925#issuecomment-2552764629
96
96
  _sym = [f"{s}.is_not_null()" for s in set(sym)]
@@ -15,7 +15,7 @@ def get_groupby_from_tuple(tup, func_name, drop_cols):
15
15
  prefix2, *_ = tup
16
16
  if len(drop_cols) > 0:
17
17
  drop_cols = [f'"{c}".into()' for c in drop_cols]
18
- drop_str = f'.drop(Selector::ByName {{ names: Arc::new([{','.join(drop_cols)}]), strict: true }})'
18
+ drop_str = f'.drop(Selector::ByName {{ names: Arc::new([{",".join(drop_cols)}]), strict: true }})'
19
19
  else:
20
20
  drop_str = ""
21
21
 
@@ -44,7 +44,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
44
44
  filename,
45
45
  date='date', asset='asset',
46
46
  extra_codes: Sequence[str] = (),
47
- over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
47
+ over_null: Literal['order_by', 'partition_by', None] = 'order_by',
48
48
  ge_date_idx: int = 0,
49
49
  **kwargs):
50
50
  """基于模板的代码生成"""
@@ -83,14 +83,14 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
83
83
  exprs_dst.append(f"#" + '=' * 40 + func_name)
84
84
  else:
85
85
  va, ex, sym, comment = kv
86
- # 多个#时,只取第一个#后的参数
87
- args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
88
86
  s1 = str(ex)
89
87
  s2 = p.doprint(ex)
90
88
  if s1 != s2:
91
89
  # 不想等,打印注释,显示会更直观察
92
90
  func_code.append(f"// {va} = {s1}")
93
91
  if k[0] == TS:
92
+ # 多个#时,只取第一个#后的参数
93
+ args, argv = parser.parse_known_args(args=comment.split("#")[1].split(" "))
94
94
  ts_func_name = func_name
95
95
  # https://github.com/pola-rs/polars/issues/12925#issuecomment-2552764629
96
96
  _sym = [f'col("{s}").is_not_null()' for s in set(sym)]
@@ -14,7 +14,7 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
14
14
  filename,
15
15
  date='date', asset='asset',
16
16
  extra_codes: Sequence[str] = (),
17
- over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
17
+ over_null: Literal['order_by', 'partition_by', None] = 'order_by',
18
18
  table_name: str = 'self',
19
19
  ge_date_idx: int = 0,
20
20
  **kwargs):
@@ -1,9 +1,10 @@
1
+ import argparse
1
2
  import inspect
2
3
  import os
3
4
  import pathlib
4
5
  from functools import lru_cache
5
6
  from io import TextIOBase
6
- from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable, Dict, Tuple
7
+ from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable
7
8
 
8
9
  import polars as pl
9
10
  from black import Mode, format_str
@@ -15,7 +16,7 @@ from sympy.simplify import cse_opts
15
16
 
16
17
  from expr_codegen.codes import sources_to_exprs
17
18
  from expr_codegen.expr import get_current_by_prefix, get_children, replace_exprs
18
- from expr_codegen.model import dag_start, dag_end, dag_middle, _RESERVED_WORD_
19
+ from expr_codegen.model import dag_start, dag_end, dag_middle, _RESERVED_WORD_, propagate_comment_recursive, create_dag_exprs
19
20
 
20
21
  try:
21
22
  from pandas import DataFrame as _pd_DataFrame
@@ -61,6 +62,20 @@ def simplify2(expr, skip_simplify: bool):
61
62
  return expr
62
63
 
63
64
 
65
+ def comment_adjust(expr_src, over_null):
66
+ parser = argparse.ArgumentParser()
67
+ parser.add_argument("--over_null", type=str, nargs="?", default=over_null)
68
+
69
+ expr_dst = []
70
+ for k, v, c in expr_src:
71
+ args, argv = parser.parse_known_args(args=c.split("#")[1].split(" "))
72
+ c = f"# --over_null={args.over_null}"
73
+
74
+ expr_dst.append((k, v, c))
75
+
76
+ return expr_dst
77
+
78
+
64
79
  class ExprTool:
65
80
 
66
81
  def __init__(self):
@@ -135,21 +150,7 @@ class ExprTool:
135
150
  syms = [str(s) for s in syms]
136
151
  return exprs, syms
137
152
 
138
- def reduce(self, repl, redu):
139
- """减少中间变量数量,有利用减少内存占用"""
140
-
141
- exprs_list = []
142
-
143
- # cse前简化一次,cse后不再简化
144
- # (~开盘涨停 & 昨收涨停) | (~收盘涨停 & 最高涨停)
145
- for k, v in repl:
146
- exprs_list.append((k, v, "#"))
147
- for k, v, c in redu:
148
- exprs_list.append((k, v, c))
149
-
150
- return exprs_list
151
-
152
- def cse(self, exprs, symbols_repl=None, exprs_src=None, skip_simplify=False):
153
+ def cse(self, exprs, symbols_repl=None, exprs_src=None, skip_simplify=False, date='date', asset='asset'):
153
154
  """多个子公式+长公式,提取公共公式
154
155
 
155
156
  Parameters
@@ -171,38 +172,26 @@ class ExprTool:
171
172
  表达式
172
173
 
173
174
  """
174
- self.exprs_names = [k for k, v, c in exprs_src]
175
- # 包含了注释信息
176
- _exprs = [k for k, v in exprs]
177
-
178
- # 注意:对于表达式右边相同,左边不同的情况,会当成一个处理
179
- if skip_simplify:
180
- repl, redu = cse(_exprs, symbols_repl, optimizations=[])
181
- else:
182
- repl, redu = cse(_exprs, symbols_repl, optimizations=[(cse_opts.sub_pre, cse_opts.sub_post), ])
183
-
184
- outputs_len = len(exprs_src)
185
-
186
- new_redu = []
187
- symbols_redu = iter(exprs_src)
188
- for expr in redu[-outputs_len:]:
189
- # 可能部分表达式只在之前出现过,后面完全用不到如,ts_rank(ts_decay_linear(x_147, 11.4157), 6.72611)
190
- variable = next(symbols_redu)
191
- a = symbols(variable[0])
192
- new_redu.append((a, expr, variable[2]))
193
-
194
- self.exprs_list = self.reduce(repl, new_redu)
195
-
196
- # with open("exprs.pickle", "wb") as file:
197
- # pickle.dump(exprs_dict, file)
198
-
175
+ self.exprs_names = [k for k, _, _ in exprs_src]
176
+ used = set(self.exprs_names) | {date, asset} | _RESERVED_WORD_
177
+ used.update(str(s) for expr, _ in exprs for s in expr.free_symbols)
178
+ pool = symbols_repl if symbols_repl is not None else numbered_symbols('_x_')
179
+ replacements = (s for s in pool if str(s) not in used)
180
+ optimizations = [] if skip_simplify else [(cse_opts.sub_pre, cse_opts.sub_post)]
181
+ self.exprs_list = []
182
+ for comment in dict.fromkeys(c for _, _, c in exprs_src):
183
+ outputs = [(k, v, c) for k, v, c in exprs_src if c == comment]
184
+ selected = [expr for expr, c in exprs if c == comment]
185
+ repl, redu = cse(selected, replacements, optimizations=optimizations)
186
+ self.exprs_list.extend((k, v, comment) for k, v in repl)
187
+ self.exprs_list.extend((symbols(k), expr, c) for (k, _, c), expr in zip(outputs, redu[-len(outputs):]))
199
188
  return self.exprs_list
200
189
 
201
- def dag(self, merge: bool, skip_columns, date, asset, skip_simplify):
190
+ def dag(self, skip_columns, date, asset):
202
191
  """生成DAG"""
203
192
  G = dag_start(self.exprs_list, self.get_current_func, self.get_current_func_kwargs, date, asset)
204
- if merge:
205
- G = dag_middle(G, self.exprs_names, skip_columns, self.get_current_func, self.get_current_func_kwargs, date, asset, skip_simplify)
193
+ G = dag_middle(G, self.exprs_names, skip_columns, self.get_current_func, self.get_current_func_kwargs, date, asset)
194
+ propagate_comment_recursive(G)
206
195
  return dag_end(G)
207
196
 
208
197
  def all(self, exprs_src, style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
@@ -210,7 +199,7 @@ class ExprTool:
210
199
  replace: bool = True, regroup: bool = False, format: bool = True,
211
200
  date='date', asset='asset',
212
201
  extra_codes: Sequence[object] = (),
213
- over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
202
+ over_null: Literal['order_by', 'partition_by', None] = 'order_by',
214
203
  table_name: str = 'self',
215
204
  ge_date_idx: int = 0,
216
205
  skip_simplify: bool = False,
@@ -254,14 +243,19 @@ class ExprTool:
254
243
  if replace:
255
244
  exprs_src = replace_exprs(exprs_src)
256
245
 
246
+ source_graph = create_dag_exprs((symbols(k), v, c) for k, v, c in exprs_src)
247
+ propagate_comment_recursive(source_graph)
248
+ exprs_context = [(k, v, source_graph.nodes[k]['comment']) for k, v, _ in exprs_src]
249
+ exprs_context = comment_adjust(exprs_context, over_null)
250
+
257
251
  # 子表达式在前,原表式在最后
258
- exprs_dst, syms_dst = self.merge(date, asset, exprs_src, skip_simplify=skip_simplify)
252
+ exprs_dst, syms_dst = self.merge(date, asset, exprs_context, skip_simplify=skip_simplify)
259
253
  syms_dst = list(set(syms_dst) - _RESERVED_WORD_)
260
254
 
261
255
  # 提取公共表达式
262
- self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), exprs_src=exprs_src, skip_simplify=skip_simplify)
256
+ self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), exprs_src=exprs_context, skip_simplify=skip_simplify, date=date, asset=asset)
263
257
  # 有向无环图流转
264
- exprs_ldl, G = self.dag(True, skip_columns, date, asset, skip_simplify=skip_simplify)
258
+ exprs_ldl, G = self.dag(skip_columns, date, asset)
265
259
 
266
260
  if regroup:
267
261
  exprs_ldl.optimize(merge=style != 'sql')
@@ -306,7 +300,7 @@ class ExprTool:
306
300
  style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
307
301
  template_file: Optional[str] = None,
308
302
  date: str = 'date', asset: str = 'asset',
309
- over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
303
+ over_null: Literal['order_by', 'partition_by', None] = 'order_by',
310
304
  table_name: str = 'self',
311
305
  ge_date_idx: int = 0,
312
306
  skip_simplify: bool = False,
@@ -1 +0,0 @@
1
- __version__ = "0.16.6"
File without changes
File without changes