expr-codegen 0.15.2__tar.gz → 0.16.1__tar.gz

{expr_codegen-0.15.2 → expr_codegen-0.16.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: expr_codegen
-Version: 0.15.2
+Version: 0.16.1
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License
@@ -276,6 +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())`，底层会试着将关键字参数转成位置参数使用
 
 ## 下划线开头的变量
 

{expr_codegen-0.15.2 → expr_codegen-0.16.1}/README.md

@@ -226,6 +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())`，底层会试着将关键字参数转成位置参数使用
 
 ## 下划线开头的变量
 

expr_codegen-0.16.1/expr_codegen/_version.py

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

{expr_codegen-0.15.2 → expr_codegen-0.16.1}/expr_codegen/codes.py

@@ -1,4 +1,5 @@
 import ast
+import inspect
 import re
 from ast import expr
 
@@ -333,6 +334,63 @@ class RenameTransformer(ast.NodeTransformer):
         return node
 
 
+class KeywordToPositionalTransformer(ast.NodeTransformer):
+    def __init__(self, function_mapping):
+        self.function_mapping = function_mapping  # 函数名到实际函数的映射
+
+    def visit_Call(self, node):
+        if isinstance(node.func, ast.Name) and node.keywords:
+            func_name = node.func.id
+            if func_name in self.function_mapping:
+                return self.transform_call(node, self.function_mapping[func_name])
+        return node
+
+    def transform_call(self, node, target_func):
+        try:
+            # 获取函数参数签名
+            sig = inspect.signature(target_func)
+            param_names = list(sig.parameters.keys())
+
+            # 构建参数映射
+            arg_mapping = {}
+            # 处理现有位置参数
+            for i, arg in enumerate(node.args):
+                if i < len(param_names):
+                    arg_mapping[param_names[i]] = arg
+
+            # 处理关键字参数
+            for keyword in node.keywords:
+                if keyword.arg in param_names:
+                    arg_mapping[keyword.arg] = keyword.value
+
+            # 按参数顺序构建新的位置参数列表
+            new_args = []
+            for param_name in param_names:
+                if param_name in arg_mapping:
+                    new_args.append(arg_mapping[param_name])
+                else:
+                    # 对于没有提供的参数，需要处理默认值
+                    param = sig.parameters[param_name]
+                    if param.default != inspect.Parameter.empty:
+                        # 使用默认值
+                        new_args.append(ast.Constant(value=param.default))
+                    else:
+                        # 必需参数缺失，保持原样或报错
+                        return node
+
+            # 创建新的调用节点
+            new_node = ast.Call(
+                func=node.func,
+                args=new_args,
+                keywords=[]
+            )
+            return new_node
+
+        except Exception as e:
+            # 转换失败时返回原节点
+            return node
+
+
 def source_replace(source: str) -> str:
     # 三元表达式转换成 错误版if( )else，一定得在Transformer中修正
     num = 1
@@ -374,7 +432,7 @@ def raw_to_code(raw):
     return '\n'.join([ast_comments.unparse(a) for a in raw])
 
 
-def sources_to_asts(*sources, convert_xor: bool):
+def sources_to_asts(*sources, convert_xor: bool, function_mapping):
     """输入多份源代码"""
 
     def _source_to_asts(source):
@@ -394,6 +452,8 @@ def sources_to_asts(*sources, convert_xor: bool):
 
     t1 = SyntaxTransformer(convert_xor)
     t1.visit(tree)
+    t2 = KeywordToPositionalTransformer(function_mapping)
+    t2.visit(tree)
     t = RenameTransformer({}, {})
     t.visit(tree)
 
@@ -429,12 +489,12 @@ def _add_default_type(globals_):
     return globals_
 
 
-def sources_to_exprs(globals_, *sources, convert_xor: bool):
+def sources_to_exprs(globals_, *sources, convert_xor: bool, function_mapping):
     """将源代码转换成表达式"""
 
     globals_ = _add_default_type(globals_)
 
-    raw, assigns, funcs_new, args_new, targets_new = sources_to_asts(*sources, convert_xor=convert_xor)
+    raw, assigns, funcs_new, args_new, targets_new = sources_to_asts(*sources, convert_xor=convert_xor, function_mapping=function_mapping)
     # 支持OPEN[1]转ts_delay(OPEN,1)
     funcs_new.add('ts_delay')
 

{expr_codegen-0.15.2 → expr_codegen-0.16.1}/expr_codegen/pandas/code.py

@@ -13,18 +13,23 @@ def get_groupby_from_tuple(tup, func_name, drop_cols):
     """从传入的元组中生成分组运行代码"""
     prefix2, *_ = tup
 
+    if len(drop_cols)>0:
+        drop_str = f'.drop(columns={drop_cols})'
+    else:
+        drop_str = ""
+
     if prefix2 == TS:
         # 组内需要按时间进行排序，需要维持顺序
         prefix2, asset = tup
-        return f'df = df.groupby(by=[_ASSET_], group_keys=False).apply({func_name}).drop(columns={drop_cols})'
+        return f'df = df.groupby(by=[_ASSET_], group_keys=False).apply({func_name}){drop_str}'
     if prefix2 == CS:
         prefix2, date = tup
-        return f'df = df.groupby(by=[_DATE_], group_keys=False).apply({func_name}).drop(columns={drop_cols})'
+        return f'df = df.groupby(by=[_DATE_], group_keys=False).apply({func_name}){drop_str}'
     if prefix2 == GP:
         prefix2, date, group = tup
-        return f'df = df.groupby(by=[_DATE_, "{group}"], group_keys=False).apply({func_name}).drop(columns={drop_cols})'
+        return f'df = df.groupby(by=[_DATE_, "{group}"], group_keys=False).apply({func_name}){drop_str}'
 
-    return f'df = {func_name}(df).drop(columns={drop_cols})'
+    return f'df = {func_name}(df){drop_str}'
 
 
 def symbols_to_code(syms):

{expr_codegen-0.15.2 → expr_codegen-0.16.1}/expr_codegen/polars/code.py

@@ -14,18 +14,23 @@ def get_groupby_from_tuple(tup, func_name, drop_cols):
     """从传入的元组中生成分组运行代码"""
     prefix2, *_ = tup
 
+    if len(drop_cols)>0:
+        drop_str = f'.drop(*{drop_cols})'
+    else:
+        drop_str = ""
+
     if prefix2 == TS:
         # 组内需要按时间进行排序，需要维持顺序
         prefix2, asset = tup
-        return f'df = {func_name}(df.sort(_ASSET_, _DATE_)).drop(*{drop_cols})'
+        return f'df = {func_name}(df.sort(_ASSET_, _DATE_)){drop_str}'
     if prefix2 == CS:
         prefix2, date = tup
-        return f'df = {func_name}(df.sort(_DATE_)).drop(*{drop_cols})'
+        return f'df = {func_name}(df.sort(_DATE_)){drop_str}'
     if prefix2 == GP:
         prefix2, date, group = tup
-        return f'df = {func_name}(df.sort(_DATE_, "{group}")).drop(*{drop_cols})'
+        return f'df = {func_name}(df.sort(_DATE_, "{group}")){drop_str}'
 
-    return f'df = {func_name}(df).drop(*{drop_cols})'
+    return f'df = {func_name}(df){drop_str}'
 
 
 def symbols_to_code(syms):

{expr_codegen-0.15.2 → expr_codegen-0.16.1}/expr_codegen/polars/template.py.j2

@@ -63,12 +63,11 @@ def {{ key }}(df: DataFrame) -> DataFrame:
 
 
 def _filter_last(df: DataFrame, ge_date_idx: int) -> DataFrame:
-    """过滤数据，只取最后几天。实盘时可用于减少计算量
-    """
+    """过滤数据，只取最后几天。实盘时可用于减少计算量"""
     if ge_date_idx == 0:
         return df
     else:
-        return df.filter(pl.col(_DATE_) >= df.select(pl.col(_DATE_).unique().sort())[ge_date_idx, 0])
+        return df.filter(pl.col(_DATE_) >= pl.col(_DATE_).unique().sort().slice(ge_date_idx, 1).first())
 
 
 def main(df: DataFrame, ge_date_idx: int) -> DataFrame:

expr_codegen-0.16.1/expr_codegen/rust/code.py

@@ -0,0 +1,150 @@
+import argparse
+import json
+import os
+from typing import Sequence, Literal
+
+import jinja2
+from jinja2 import FileSystemLoader, TemplateNotFound
+
+from expr_codegen.expr import TS, CS, GP
+from expr_codegen.model import ListDictList
+from expr_codegen.rust.printer import RustStrPrinter
+
+
+def get_groupby_from_tuple(tup, func_name, drop_cols):
+    """从传入的元组中生成分组运行代码"""
+    prefix2, *_ = tup
+    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:
+        drop_str = ""
+
+    if prefix2 == TS:
+        # 组内需要按时间进行排序，需要维持顺序
+        prefix2, asset = tup
+        return f'df = {func_name}(df.sort([_ASSET_, _DATE_], SortMultipleOptions::default())){drop_str};'
+    if prefix2 == CS:
+        prefix2, date = tup
+        return f'df = {func_name}(df.sort([_DATE_], SortMultipleOptions::default())){drop_str};'
+    if prefix2 == GP:
+        prefix2, date, group = tup
+        return f'df = {func_name}(df.sort([_DATE_, "{group}"], SortMultipleOptions::default())){drop_str};'
+
+    return f'df = {func_name}(df){drop_str};'
+
+
+# def symbols_to_code(syms):
+#     a = [f"{s}" for s in syms]
+#     b = [f"'{s}'" for s in syms]
+#     return f"""_ = [{','.join(b)}]
+# [{','.join(a)}] = [pl.col(i) for i in _]"""
+
+
+def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
+            filename,
+            date='date', asset='asset',
+            extra_codes: Sequence[str] = (),
+            over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
+            ge_date_idx: int = 0,
+            **kwargs):
+    """基于模板的代码生成"""
+    if filename is None:
+        filename = 'template.rs.j2'
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--over_null", type=str, nargs="?", default=over_null)
+
+    # 打印Polars风格代码
+    p = RustStrPrinter()
+
+    # polars风格代码
+    funcs = {}
+    # 分组应用代码。这里利用了字典按插入顺序排序的特点，将排序放在最前
+    groupbys = {'sort': ''}
+    # 处理过后的表达式
+    exprs_dst = []
+    syms_out = []
+    ts_func_name = None
+    drop_symbols = exprs_ldl.drop_symbols()
+    j = -1
+    for i, row in enumerate(exprs_ldl.values()):
+        for k, vv in row.items():
+            j += 1
+            if len(vv) == 0:
+                continue
+            # 函数名
+            func_name = f'func_{i}_{"_".join(k)}'
+            func_code = []
+            for kv in vv:
+                if kv is None:
+                    func_code.append(f"    ]);")
+                    func_code.append(f"// " + '=' * 40)
+                    func_code.append(f"    df = df.with_columns([")
+                    exprs_dst.append(f"#" + '=' * 40 + func_name)
+                else:
+                    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:
+                        # 不想等，打印注释，显示会更直观察
+                        func_code.append(f"// {va} = {s1}")
+                    if k[0] == TS:
+                        ts_func_name = func_name
+                        # https://github.com/pola-rs/polars/issues/12925#issuecomment-2552764629
+                        _sym = [f'col("{s}").is_not_null()' for s in set(sym)]
+                        if len(_sym) > 1:
+                            _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:
+                            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:
+                        func_code.append(f'({s2}).over([_DATE_, "{k[2]}"]).alias("{va}"),')
+                    else:
+                        func_code.append(f'({s2}).alias("{va}"),')
+                    exprs_dst.append(f"{va} = {s1} {comment}")
+                    if va not in syms_dst:
+                        syms_out.append(va)
+            func_code.append(f"    ]);")
+            func_code = func_code[1:]
+
+            # polars风格代码列表
+            funcs[func_name] = '\n'.join(func_code)
+            # 只有下划线开头再删除
+            ds = [x for x in drop_symbols[j] if x.startswith('_')]
+            # 分组应用代码
+            groupbys[func_name] = get_groupby_from_tuple(k, func_name, ds)
+
+    # syms1 = symbols_to_code(syms_dst)
+    # syms2 = symbols_to_code(syms_out)
+    # filter_last处理
+    _groupbys = {'sort': groupbys['sort']}
+    if ts_func_name is None:
+        _groupbys['_filter_last'] = "df = _filter_last(df, ge_date_idx);"
+    for k, v in groupbys.items():
+        _groupbys[k] = v
+        if k == ts_func_name:
+            _groupbys[k + '_filter_last'] = "df = _filter_last(df, ge_date_idx);"
+    groupbys = _groupbys
+
+    try:
+        env = jinja2.Environment(loader=FileSystemLoader(os.path.dirname(__file__)))
+        template = env.get_template(filename)
+    except TemplateNotFound:
+        env = jinja2.Environment(loader=FileSystemLoader(os.path.dirname(filename)))
+        template = env.get_template(os.path.basename(filename))
+
+    return template.render(funcs=funcs, groupbys=groupbys,
+                           exprs_src=exprs_src, exprs_dst=exprs_dst,
+                           # syms1=syms1, syms2=syms2,
+                           date=date, asset=asset,
+                           extra_codes=extra_codes)

expr_codegen-0.16.1/expr_codegen/rust/printer.py

@@ -0,0 +1,115 @@
+import inspect
+
+from sympy import Basic, Function, StrPrinter
+
+
+# TODO: 如有新添加函数，但表达式有变更才需要在此补充对应的打印代码，否则可以省略
+
+class RustStrPrinter(StrPrinter):
+    def _print(self, expr, **kwargs) -> str:
+        """Internal dispatcher
+
+        Tries the following concepts to print an expression:
+            1. Let the object print itself if it knows how.
+            2. Take the best fitting method defined in the printer.
+            3. As fall-back use the emptyPrinter method for the printer.
+        """
+        self._print_level += 1
+        try:
+            # If the printer defines a name for a printing method
+            # (Printer.printmethod) and the object knows for itself how it
+            # should be printed, use that method.
+            if self.printmethod and hasattr(expr, self.printmethod):
+                if not (isinstance(expr, type) and issubclass(expr, Basic)):
+                    return getattr(expr, self.printmethod)(self, **kwargs)
+
+            # See if the class of expr is known, or if one of its super
+            # classes is known, and use that print function
+            # Exception: ignore the subclasses of Undefined, so that, e.g.,
+            # Function('gamma') does not get dispatched to _print_gamma
+            classes = type(expr).__mro__
+            # if AppliedUndef in classes:
+            #     classes = classes[classes.index(AppliedUndef):]
+            # if UndefinedFunction in classes:
+            #     classes = classes[classes.index(UndefinedFunction):]
+            # Another exception: if someone subclasses a known function, e.g.,
+            # gamma, and changes the name, then ignore _print_gamma
+            if Function in classes:
+                i = classes.index(Function)
+                classes = tuple(c for c in classes[:i] if \
+                                c.__name__ == classes[0].__name__ or \
+                                c.__name__.endswith("Base")) + classes[i:]
+            for cls in classes:
+                printmethodname = '_print_' + cls.__name__
+
+                # 所有以gp_开头的函数都转换成cs_开头
+                if printmethodname.startswith('_print_gp_'):
+                    printmethodname = "_print_gp_"
+
+                printmethod = getattr(self, printmethodname, None)
+                if printmethod is not None:
+                    return printmethod(expr, **kwargs)
+            # Unknown object, fall back to the emptyPrinter.
+            return self.emptyPrinter(expr)
+        finally:
+            self._print_level -= 1
+
+    def _print_Symbol(self, expr):
+        if expr.name in ('_NONE_', '_TRUE_', '_FALSE_'):
+            return expr.name
+        return f'col("{expr.name}")'
+
+    def _print_Equality(self, expr):
+        new_args = [f"eq({self._print(arg)})" for arg in expr.args]
+        return ".".join(new_args)[2:]
+
+    def _print_Or(self, expr):
+        new_args = [f"or({self._print(arg)})" for arg in expr.args]
+        return ".".join(new_args)[2:]
+
+    def _print_Xor(self, expr):
+        new_args = [f"xor({self._print(arg)})" for arg in expr.args]
+        return ".".join(new_args)[3:]
+
+    def _print_And(self, expr):
+        new_args = [f"and({self._print(arg)})" for arg in expr.args]
+        return ".".join(new_args)[3:]
+
+    def _print_Not(self, expr):
+        return "(%s).not()" % self._print(expr.args[0])
+
+    def _print_gp_(self, expr):
+        """gp_函数都转换成cs_函数，但要丢弃第一个参数"""
+        new_args = [self._print(arg) for arg in expr.args[1:]]
+        func_name = expr.func.__name__[3:]
+        return "cs_%s(%s)" % (func_name, ",".join(new_args))
+
+    def _print_Integer(self, expr):
+        caller_frame = inspect.stack()[2]
+        caller_name = caller_frame.function
+        if caller_name in ("_print_Pow", "_print_Add", "_print_Mul", "_print_Relational"):
+            return "lit(%s)" % super()._print_Integer(expr)
+        else:
+            return super()._print_Integer(expr)
+
+    def _print_Float(self, expr):
+        caller_frame = inspect.stack()[2]
+        caller_name = caller_frame.function
+        if caller_name in ("_print_Pow", "_print_Add", "_print_Mul", "_print_Relational"):
+            return "lit(%s)" % super()._print_Float(expr)
+        else:
+            return super()._print_Float(expr)
+
+    def _print_Relational(self, expr):
+
+        charmap = {
+            "<": "lt",
+            ">": "gt",
+            ">=": "gt_eq",
+            "<=": "lt_eq",
+        }
+
+        if expr.rel_op in charmap:
+            return '(%s).%s(%s)' % (self._print(expr.lhs), charmap[expr.rel_op], self._print(expr.rhs))
+
+        return super()._print_Relational(expr)

expr_codegen-0.16.1/expr_codegen/rust/template.rs.j2

@@ -0,0 +1,125 @@
+// this code is auto generated by the expr_codegen
+// https://github.com/wukan1986/expr_codegen
+// 此段代码由 expr_codegen 自动生成，欢迎提交 issue 或 pull request
+use polars::prelude::*;
+
+pub fn ts_delay(x: Expr, n: i16) -> Expr {
+    x.shift(lit(n))
+}
+
+pub fn ts_mean(x: Expr, d: usize) -> Expr {
+    x.rolling_mean(RollingOptionsFixedWindow {
+        window_size: d,
+        min_periods: d,
+        weights: None,
+        center: false,
+        fn_params: None,
+    })
+}
+
+pub fn ts_sum(x: Expr, d: usize) -> Expr {
+    x.rolling_sum(RollingOptionsFixedWindow {
+        window_size: d,
+        min_periods: d,
+        weights: None,
+        center: false,
+        fn_params: None,
+    })
+}
+
+pub fn cs_zscore(x: Expr, ddof: u8) -> Expr {
+    (x.clone() - x.clone().mean()) / x.clone().std(ddof)
+}
+
+const _DATE_: &str = "date";
+const _ASSET_: &str = "asset";
+const _NONE_: Option<i32> = None;
+const _TRUE_: bool = true;
+const _FALSE_: bool = false;
+
+
+
+{%-for row in extra_codes %}
+{{ row-}}
+{% endfor %}
+
+{% for key, value in funcs.items() %}
+
+fn {{ key }}(mut df: LazyFrame) -> LazyFrame{
+{{ value }}
+    df}
+
+{% endfor %}
+
+/*
+{%-for row in exprs_dst %}
+{{ row-}}
+{% endfor %}
+*/
+
+/*
+{%-for a,b,c in exprs_src %}
+{{ a }} = {{ b}} {{c-}}
+{% endfor %}
+*/
+
+pub fn main_(mut df: LazyFrame, ge_date_idx: i64) -> LazyFrame {
+    {% for key, value in groupbys.items() %}
+    {{ value-}}
+    {% endfor %}
+
+    df
+}
+
+pub fn _filter_last(df: LazyFrame, ge_date_idx: i64) -> LazyFrame {
+    if ge_date_idx == 0 {
+        df
+    } else {
+        let date_expr = col(_DATE_)
+            .unique()
+            .sort(SortOptions::default())
+            .slice(ge_date_idx, 1)
+            .first();
+
+        df.filter(col(_DATE_).gt_eq(date_expr))
+    }
+}
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // 目前生成的代码还需要调整才能投入使用，还有更多的函数需要补充
+    let mut df = df! (
+        "date" => [1, 2, 3, 4, 5, 1, 2, 3, 4],
+        "asset" => [1, 2, 3, 1, 2, 3, 1, 2, 3],
+        "OPEN" => [1, 2, 3, 4, 5, 6, 7, 8, 9],
+        "HIGH" => [1, 2, 3, 4, 5, 6, 7, 8, 9],
+        "LOW" => [1, 2, 3, 4, 5, 6, 7, 8, 9],
+        "CLOSE" => [1, 2, 3, 4, 5, 6, 7, 8, 9],
+    )?;
+
+    df = main_(df.lazy(), 0)
+        .drop(Selector::Matches("^_.*$".into()))
+        .collect()?;
+    println!("{:?}", df);
+
+    Ok(())
+}
+
+/*
+# Cargo.toml
+# https://docs.pola.rs/user-guide/installation/
+
+[package]
+name = "expr_codegen_rs"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+polars = { version = "0.51.0", features = [
+    "lazy",
+    "round_series",
+    "strings",
+    "regex",
+    "rolling_window",
+] }
+
+*/

{expr_codegen-0.15.2 → expr_codegen-0.16.1}/expr_codegen/tool.py

@@ -2,7 +2,7 @@ import inspect
 import pathlib
 from functools import lru_cache
 from io import TextIOBase
-from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable
+from typing import Sequence, Union, TypeVar, Optional, Literal, Iterable, Dict
 
 import polars as pl
 from black import Mode, format_str
@@ -201,7 +201,7 @@ class ExprTool:
                            asset)
         return dag_end(G)
 
-    def all(self, exprs_src, style: Literal['pandas', 'polars', 'sql'] = 'polars',
+    def all(self, exprs_src, style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
             template_file: Optional[str] = None,
             replace: bool = True, regroup: bool = False, format: bool = True,
             date='date', asset='asset',
@@ -245,7 +245,7 @@ class ExprTool:
         代码字符串
 
         """
-        assert style in ('pandas', 'polars', 'sql')
+        assert style in ('pandas', 'polars', 'sql', 'rust')
 
         if replace:
             exprs_src = replace_exprs(exprs_src)
@@ -269,6 +269,9 @@ class ExprTool:
         elif style == 'sql':
             from expr_codegen.sql.code import codegen
             format = False
+        elif style == 'rust':
+            from expr_codegen.rust.code import codegen
+            format = False
         else:
             raise ValueError(f'unknown style {style}')
 
@@ -290,13 +293,13 @@ 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,
                   output_file: str,
                   convert_xor: bool,
-                  style: Literal['pandas', 'polars', 'sql'] = 'polars',
+                  style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
                   template_file: Optional[str] = None,
                   date: str = 'date', asset: str = 'asset',
                   over_null: Literal['order_by', 'partition_by', None] = 'partition_by',
@@ -304,9 +307,10 @@ 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)
+        raw, exprs_list = sources_to_exprs(self.globals_, source, *more_sources, convert_xor=convert_xor, function_mapping=function_mapping)
 
         # 生成代码
         code, G = _TOOL_.all(exprs_list, style=style, template_file=template_file,
@@ -380,13 +384,14 @@ def codegen_exec(df: Union[DataFrame, None],
                  output_file: Union[str, TextIOBase, None] = None,
                  run_file: Union[bool, str] = False,
                  convert_xor: bool = False,
-                 style: Literal['pandas', 'polars', 'sql'] = 'polars',
+                 style: Literal['pandas', 'polars', 'sql', 'rust'] = 'polars',
                  template_file: Optional[str] = None,
                  date: str = 'date', asset: str = 'asset',
                  table_name: str = 'self',
                  ge_date_idx: int = 0,
                  skip_simplify: bool = False,
                  skip_columns: Iterable[str] = (),
+                 function_mapping: Dict = {},
                  **kwargs) -> Union[DataFrame, str]:
     """快速转换源代码并执行
 
@@ -437,6 +442,8 @@ def codegen_exec(df: Union[DataFrame, None],
         已经存在的列不参与计算。可用于加快计算速度。只在计算耗时久时再用，否则没有必要
         例如：在研发阶段，第一次计算100个因子，第二次，只改动了其中的5个，所以只要将这5个从df.columns中排除即可。
         注意：生成的源代码有差异。
+    function_mapping:
+        传入函数定义，可直接传`globals()`。用于将所有的关键字参数转换成位置参数
 
     Returns
     -------
@@ -494,12 +501,15 @@ 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
     )
 
     if df is None:
         # 如果df为空，直接返回代码
         return code
+    elif style == 'rust':
+        return code
     elif style == 'sql':
         with pl.SQLContext(frames={table_name: df}) as ctx:
             return ctx.execute(code, eager=isinstance(df, _pl_DataFrame))

expr_codegen-0.15.2/expr_codegen/_version.py

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