expr-codegen 0.16.2__tar.gz → 0.16.4__tar.gz

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: expr_codegen
-Version: 0.16.2
+Version: 0.16.4
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License
@@ -276,7 +276,7 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
    ```
 11. 单行注释支持参数输入，如：`# --over_null`、`# --over_null=order_by`、`# --over_null=partition_by`
 12. 代码块中，对`import`、`def`、`class`三种语法，自动原样插入到生成的代码中
-13. 由于`sympy`的限制不支持关键字参数，但如果`from polars_ta.prefix.wq import *`,然后`codegen_exec(function_mapping=globals())`，底层会试着将关键字参数转成位置参数使用
+13. 虽然`sympy`的限制不支持关键字参数，但`codegen_exec`底层会试着将关键字参数转成位置参数使用
 
 ## 下划线开头的变量
 

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/README.md

@@ -226,7 +226,7 @@ X3 = (ts_returns(CLOSE, 3)).over(_ASSET_, order_by=_DATE_),
    ```
 11. 单行注释支持参数输入，如：`# --over_null`、`# --over_null=order_by`、`# --over_null=partition_by`
 12. 代码块中，对`import`、`def`、`class`三种语法，自动原样插入到生成的代码中
-13. 由于`sympy`的限制不支持关键字参数，但如果`from polars_ta.prefix.wq import *`,然后`codegen_exec(function_mapping=globals())`，底层会试着将关键字参数转成位置参数使用
+13. 虽然`sympy`的限制不支持关键字参数，但`codegen_exec`底层会试着将关键字参数转成位置参数使用
 
 ## 下划线开头的变量
 

expr_codegen-0.16.4/expr_codegen/_version.py

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

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/expr_codegen/codes.py

@@ -489,9 +489,9 @@ def _add_default_type(globals_):
     return globals_
 
 
-def sources_to_exprs(globals_, *sources, convert_xor: bool, function_mapping):
+def sources_to_exprs(globals_, *sources, convert_xor: bool):
     """将源代码转换成表达式"""
-
+    function_mapping = {k: v for k, v in globals_.items() if inspect.isfunction(v)}
     globals_ = _add_default_type(globals_)
 
     raw, assigns, funcs_new, args_new, targets_new = sources_to_asts(*sources, convert_xor=convert_xor, function_mapping=function_mapping)

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/expr_codegen/model.py

@@ -365,12 +365,14 @@ def dag_start(exprs_list, func, func_kwargs, date, asset):
     return G
 
 
-def dag_middle(G, exprs_names, skip_columns, func, func_kwargs, date, asset):
+def dag_middle(G, exprs_names, skip_columns, func, func_kwargs, date, asset, skip_simplify):
     """删除几个没有必要的节点"""
     # 以下划线开头的节点，不保留
     keep_nodes = [k for k in exprs_names if not k.startswith('_')]
 
-    G = merge_nodes_1(G, keep_nodes, *keep_nodes)
+    if not skip_simplify:
+        # ts_rank(-RET - -RET, 20),防止替换成ts_rank(0, 20)
+        G = merge_nodes_1(G, keep_nodes, *keep_nodes)
     G = merge_nodes_2(G, keep_nodes, *keep_nodes)
 
     # 移除0出度的节点，但保留部分

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/expr_codegen/polars/code.py

@@ -14,7 +14,7 @@ def get_groupby_from_tuple(tup, func_name, drop_cols):
     """从传入的元组中生成分组运行代码"""
     prefix2, *_ = tup
 
-    if len(drop_cols)>0:
+    if len(drop_cols) > 0:
         drop_str = f'.drop(*{drop_cols})'
     else:
         drop_str = ""
@@ -98,12 +98,17 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                             _sym = f"pl.all_horizontal({','.join(_sym)})"
                         else:
                             _sym = ','.join(_sym)
-                        if args.over_null == 'partition_by':
-                            func_code.append(f"{va}=({s2}).over({_sym}, _ASSET_, order_by=_DATE_),")
-                        elif args.over_null == 'order_by':
-                            func_code.append(f"{va}=({s2}).over(_ASSET_, order_by=[{_sym}, _DATE_]),")
-                        else:
+
+                        if len(_sym) == 0:
                             func_code.append(f"{va}=({s2}).over(_ASSET_, order_by=_DATE_),")
+                        else:
+                            if args.over_null == 'partition_by':
+                                func_code.append(f"{va}=({s2}).over({_sym}, _ASSET_, order_by=_DATE_),")
+                            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_),")
+
                     elif k[0] == CS:
                         func_code.append(f"{va}=({s2}).over(_DATE_),")
                     elif k[0] == GP:

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/expr_codegen/polars/template.py.j2

@@ -19,6 +19,7 @@ from polars_ta.prefix.ta import *  # noqa
 from polars_ta.prefix.wq import *  # noqa
 from polars_ta.prefix.cdl import *  # noqa
 from polars_ta.prefix.vec import *  # noqa
+from polars_ta.utils.functions import apply_const_to_expr  # noqa
 
 DataFrame = TypeVar('DataFrame', _pl_LazyFrame, _pl_DataFrame)
 # ===================================

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/expr_codegen/rust/code.py

@@ -1,5 +1,4 @@
 import argparse
-import json
 import os
 from typing import Sequence, Literal
 
@@ -14,7 +13,7 @@ from expr_codegen.rust.printer import RustStrPrinter
 def get_groupby_from_tuple(tup, func_name, drop_cols):
     """从传入的元组中生成分组运行代码"""
     prefix2, *_ = tup
-    if len(drop_cols)>0:
+    if len(drop_cols) > 0:
         drop_cols = [f'"{c}".into()' for c in drop_cols]
         drop_str = f'.drop(Selector::ByName {{ names: Arc::new([{','.join(drop_cols)}]), strict: true }})'
     else:
@@ -99,12 +98,17 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                             _sym = f"all_horizontal([{','.join(_sym)}]).unwrap()"
                         else:
                             _sym = ','.join(_sym)
-                        if args.over_null == 'partition_by':
-                            func_code.append(f'({s2}).over_with_options(Some([{_sym}, col(_ASSET_)]), Some(([col(_DATE_), lit(1)], SortOptions::default())), WindowMapping::default()).unwrap().alias("{va}"),')
-                        elif args.over_null == 'order_by':
-                            func_code.append(f'({s2}).over_with_options(Some([col(_ASSET_), lit(1)]), Some(([{_sym}, col(_DATE_)], SortOptions::default())), WindowMapping::default()).unwrap().alias("{va}"),')
-                        else:
+
+                        if len(_sym) == 0:
                             func_code.append(f'({s2}).over_with_options(Some([_ASSET_]), Some(([_DATE_], SortOptions::default())), WindowMapping::default()).unwrap().alias("{va}"),')
+                        else:
+                            if args.over_null == 'partition_by':
+                                func_code.append(f'({s2}).over_with_options(Some([{_sym}, col(_ASSET_)]), Some(([col(_DATE_), lit(1)], SortOptions::default())), WindowMapping::default()).unwrap().alias("{va}"),')
+                            elif args.over_null == 'order_by':
+                                func_code.append(f'({s2}).over_with_options(Some([col(_ASSET_), lit(1)]), Some(([{_sym}, col(_DATE_)], SortOptions::default())), WindowMapping::default()).unwrap().alias("{va}"),')
+                            else:
+                                func_code.append(f'({s2}).over_with_options(Some([_ASSET_]), Some(([_DATE_], SortOptions::default())), WindowMapping::default()).unwrap().alias("{va}"),')
+
                     elif k[0] == CS:
                         func_code.append(f'({s2}).over([_DATE_]).alias("{va}"),')
                     elif k[0] == GP:

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/expr_codegen/sql/code.py

@@ -65,12 +65,17 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                             _sym = f"({' AND '.join(_sym)})"
                         else:
                             _sym = ','.join(_sym)
-                        if args.over_null == 'partition_by':
-                            func_code.append(f"{s2} OVER(PARTITION BY {_sym},`{asset}` ORDER BY `{date}`) AS {va},")
-                        elif args.over_null == 'order_by':
-                            func_code.append(f"{s2} OVER(PARTITION BY `{asset}` ORDER BY {_sym},`{date}`) AS {va},")
-                        else:
+
+                        if len(_sym) == 0:
                             func_code.append(f"{s2} OVER(PARTITION BY `{asset}` ORDER BY `{date}`) AS {va},")
+                        else:
+                            if args.over_null == 'partition_by':
+                                func_code.append(f"{s2} OVER(PARTITION BY {_sym},`{asset}` ORDER BY `{date}`) AS {va},")
+                            elif args.over_null == 'order_by':
+                                func_code.append(f"{s2} OVER(PARTITION BY `{asset}` ORDER BY {_sym},`{date}`) AS {va},")
+                            else:
+                                func_code.append(f"{s2} OVER(PARTITION BY `{asset}` ORDER BY `{date}`) AS {va},")
+
                     elif k[0] == CS:
                         func_code.append(f"{s2} OVER(PARTITION BY `{date}`) AS {va},")
                     elif k[0] == GP:

{expr_codegen-0.16.2 → expr_codegen-0.16.4}/expr_codegen/tool.py

@@ -1,8 +1,9 @@
 import inspect
+import os
 import pathlib
 from functools import lru_cache
 from io import TextIOBase
-from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable, Dict
+from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable, Dict, Tuple
 
 import polars as pl
 from black import Mode, format_str
@@ -148,7 +149,7 @@ class ExprTool:
 
         return exprs_list
 
-    def cse(self, exprs, symbols_repl=None, exprs_src=None):
+    def cse(self, exprs, symbols_repl=None, exprs_src=None, skip_simplify=False):
         """多个子公式+长公式，提取公共公式
 
         Parameters
@@ -175,7 +176,11 @@ class ExprTool:
         _exprs = [k for k, v in exprs]
 
         # 注意：对于表达式右边相同，左边不同的情况，会当成一个处理
-        repl, redu = cse(_exprs, symbols_repl, optimizations=[(cse_opts.sub_pre, cse_opts.sub_post), ])
+        if skip_simplify:
+            repl, redu = cse(_exprs, symbols_repl, optimizations=[])
+        else:
+            repl, redu = cse(_exprs, symbols_repl, optimizations=[(cse_opts.sub_pre, cse_opts.sub_post), ])
+
         outputs_len = len(exprs_src)
 
         new_redu = []
@@ -193,12 +198,11 @@ class ExprTool:
 
         return self.exprs_list
 
-    def dag(self, merge: bool, skip_columns, date, asset):
+    def dag(self, merge: bool, skip_columns, date, asset, skip_simplify):
         """生成DAG"""
         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, skip_columns, self.get_current_func, self.get_current_func_kwargs, date,
-                           asset)
+            G = dag_middle(G, self.exprs_names, skip_columns, self.get_current_func, self.get_current_func_kwargs, date, asset, skip_simplify)
         return dag_end(G)
 
     def all(self, exprs_src, style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
@@ -251,13 +255,13 @@ class ExprTool:
             exprs_src = replace_exprs(exprs_src)
 
         # 子表达式在前，原表式在最后
-        exprs_dst, syms_dst = self.merge(date, asset, exprs_src, skip_simplify)
+        exprs_dst, syms_dst = self.merge(date, asset, exprs_src, skip_simplify=skip_simplify)
         syms_dst = list(set(syms_dst) - _RESERVED_WORD_)
 
         # 提取公共表达式
-        self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), exprs_src=exprs_src)
+        self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), exprs_src=exprs_src, skip_simplify=skip_simplify)
         # 有向无环图流转
-        exprs_ldl, G = self.dag(True, skip_columns, date, asset)
+        exprs_ldl, G = self.dag(True, skip_columns, date, asset, skip_simplify=skip_simplify)
 
         if regroup:
             exprs_ldl.optimize(merge=style != 'sql')
@@ -293,7 +297,7 @@ class ExprTool:
 
         return codes, G
 
-    # @lru_cache(maxsize=64)
+    @lru_cache(maxsize=64)
     def _get_code(self,
                   source: str, *more_sources: str,
                   extra_codes: str,
@@ -307,10 +311,9 @@ class ExprTool:
                   ge_date_idx: int = 0,
                   skip_simplify: bool = False,
                   skip_columns: Iterable[str] = (),
-                  function_mapping={},
                   **kwargs) -> str:
         """通过字符串生成代码， 加了缓存，多次调用不重复生成"""
-        raw, exprs_list = sources_to_exprs(self.globals_, source, *more_sources, convert_xor=convert_xor, function_mapping=function_mapping)
+        raw, exprs_list = sources_to_exprs(self.globals_, source, *more_sources, convert_xor=convert_xor)
 
         # 生成代码
         code, G = _TOOL_.all(exprs_list, style=style, template_file=template_file,
@@ -391,9 +394,8 @@ def codegen_exec(df: Union[DataFrame, None],
                  ge_date_idx: int = 0,
                  skip_simplify: bool = False,
                  skip_columns: Iterable[str] = (),
-                 function_mapping: Dict = {},
                  **kwargs) -> Union[DataFrame, str]:
-    """快速转换源代码并执行
+    r"""快速转换源代码并执行
 
     Parameters
     ----------
@@ -403,7 +405,12 @@ def codegen_exec(df: Union[DataFrame, None],
     codes:
         函数体。此部分中的表达式会被翻译成目标代码
     extra_codes: str
-        额外代码。不做处理，会被直接复制到目标代码中
+        额外代码。不做处理，会被直接复制到目标代码中。例如：
+
+        r'CS_SW_L1 = r"^sw_l1_\d+$"'
+
+        apply_const_to_expr()
+
     output_file: str| TextIOBase
         保存生成的目标代码到文件中
     run_file: bool or str
@@ -412,6 +419,7 @@ def codegen_exec(df: Union[DataFrame, None],
         - 如果是字符串，会自动从run_file中读取代码
         - 如果是模块名，会自动从模块中读取代码(可调试)
             - 注意：可能调用到其他目录下的同名模块，所以保存的文件名要有辨识度
+        - 如果文件不存在，先生成文件。第二次从生成的文件中生成
     convert_xor: bool
         ^ 转成异或还是乘方
     style: str
@@ -438,12 +446,13 @@ def codegen_exec(df: Union[DataFrame, None],
         -2 表示最近两天 >=date[-2]
     skip_simplify:bool
         遗传算法时很有可能出现OPEN/OPEN，可以跳过化简步骤
+        1. 跳过cse前的simplify
+        2. 跳过cse时的optimizations
+        3. 跳过DAG中的部分merge步骤
     skip_columns:
         已经存在的列不参与计算。可用于加快计算速度。只在计算耗时久时再用，否则没有必要
         例如：在研发阶段，第一次计算100个因子，第二次，只改动了其中的5个，所以只要将这5个从df.columns中排除即可。
         注意：生成的源代码有差异。
-    function_mapping:
-        传入函数定义，可直接传`globals()`。用于将所有的关键字参数转换成位置参数
 
     Returns
     -------
@@ -473,10 +482,12 @@ def codegen_exec(df: Union[DataFrame, None],
 
         if input_file is not None:
             if input_file.endswith('.py'):
-                return _get_func_from_file_py(input_file)(df, ge_date_idx)
+                if os.path.exists(input_file):
+                    return _get_func_from_file_py(input_file)(df, ge_date_idx)
             elif input_file.endswith('.sql'):
-                with pl.SQLContext(frames={table_name: df}) as ctx:
-                    return ctx.execute(_get_code_from_file(input_file), eager=isinstance(df, _pl_DataFrame))
+                if os.path.exists(input_file):
+                    with pl.SQLContext(frames={table_name: df}) as ctx:
+                        return ctx.execute(_get_code_from_file(input_file), eager=isinstance(df, _pl_DataFrame))
             else:
                 return _get_func_from_module(input_file)(df, ge_date_idx)  # 可断点调试
     else:
@@ -501,7 +512,6 @@ def codegen_exec(df: Union[DataFrame, None],
         ge_date_idx=ge_date_idx,
         skip_simplify=skip_simplify,
         skip_columns=skip_columns,
-        function_mapping=function_mapping,
         **kwargs
     )
 

expr_codegen-0.16.2/expr_codegen/_version.py

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