expr-codegen 0.8.7__tar.gz → 0.9.1__tar.gz

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: expr_codegen
-Version: 0.8.7
+Version: 0.9.1
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License
@@ -206,6 +206,7 @@ df = codegen_exec(df, _code_block_1, _code_block_2)  # 只执行，不保存代
 6. 不支持`A<=B<=C`，需手工替换成`(A<=B)&(B<=C)`
 7. 支持`A[0]+B[1]+C[2]`，底层会转成`A+ts_delay(B,1)+ts_delay(C,2)`
 8. 支持`~A`,底层会转换成`Not(A)`
+9. `gp_`开头的函数都会返回对应的`cs_`函数。如`gp_func(A,B,C)`会替换成`cs_func(B,C)`,其中`A`用在了`groupby([date, A])`
 
 ## 下划线开头的变量
 

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/README.md

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

expr_codegen-0.9.1/expr_codegen/_version.py

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

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/expr_codegen/codes.py

@@ -10,8 +10,9 @@ from expr_codegen.expr import register_symbols, dict_to_exprs
 class SyntaxTransformer(ast.NodeTransformer):
     """修改语法。注意：一定要修改语法后才能改名"""
 
-    # ^ 是异或还是乘方呢？
-    convert_xor: bool = True
+    def __init__(self, convert_xor):
+        # ^ 是异或还是乘方呢？
+        self.convert_xor = convert_xor
 
     def visit_Compare(self, node):
         assert len(node.comparators) == 1, f"不支持连续等号，请手工添加括号, {ast.unparse(node)}"
@@ -24,7 +25,7 @@ class SyntaxTransformer(ast.NodeTransformer):
         # 只要body区域，出现了or True，就认为是特殊处理过的
         if isinstance(node.body, ast.BoolOp) and isinstance(node.body.op, ast.Or):
             if isinstance(node.body.values[-1], ast.Constant):
-                if node.body.values[-1].value == True:
+                if node.body.values[-1].value:
                     node.test, node.body = node.body.values[0], node.test
 
         node = ast.Call(
@@ -87,7 +88,9 @@ class SyntaxTransformer(ast.NodeTransformer):
         return node
 
     def visit_Subscript(self, node):
-        if node.slice.value == 0:
+        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:
             node = node.value
         else:
             node = ast.Call(
@@ -102,25 +105,22 @@ class SyntaxTransformer(ast.NodeTransformer):
 class RenameTransformer(ast.NodeTransformer):
     """改名处理。改名前需要语法规范"""
 
-    # 旧记录
-    funcs_old = set()
-    args_old = set()
-    targets_old = set()
-    # 旧记录
-    funcs_new = set()
-    args_new = set()
-    targets_new = set()
-
-    # 映射
-    funcs_map = {}
-    # 由于None等常量无法在sympy中正确处理，只能改成Symbol变量
-    # !!!一定要在drop_symbols时排除
-    args_map = {'True': "_TRUE_", 'False': "_FALSE_", 'None': "_NONE_"}
-    targets_map = {}  # 只对非下划线开头的生效
-
-    def config_map(self, funcs_map, args_map, targets_map):
+    def __init__(self, funcs_map, targets_map, args_map=None):
+
+        if args_map is None:
+            args_map = {'True': "_TRUE_", 'False': "_FALSE_", 'None': "_NONE_"}
+        self.funcs_old = set()
+        self.args_old = set()
+        self.targets_old = set()
+        self.funcs_new = set()
+        self.args_new = set()
+        self.targets_new = set()
+        # 映射
         self.funcs_map = funcs_map
+        # 由于None等常量无法在sympy中正确处理，只能改成Symbol变量
+        # !!!一定要在drop_symbols时排除
         self.args_map = args_map
+        # 只对非下划线开头的生效
         self.targets_map = targets_map
 
     def visit_Call(self, node):
@@ -288,21 +288,6 @@ class RenameTransformer(ast.NodeTransformer):
         return node
 
 
-def sources_to_asts(*sources, convert_xor: bool):
-    """输入多份源代码"""
-    raw = []
-    assigns = {}
-    funcs_new, args_new, targets_new = set(), set(), set()
-    for arg in sources:
-        r, a, funcs_, args_, targets_ = _source_to_asts(arg, convert_xor)
-        raw.append(r)
-        assigns.update(a)
-        funcs_new.update(funcs_)
-        args_new.update(args_)
-        targets_new.update(targets_)
-    return '\n'.join(raw), assigns, funcs_new, args_new, targets_new
-
-
 def source_replace(source: str) -> str:
     # 三元表达式转换成 错误版if( )else，一定得在Transformer中修正
     num = 1
@@ -316,24 +301,43 @@ def source_replace(source: str) -> str:
     return source
 
 
-def _source_to_asts(source, convert_xor: bool):
-    """源代码"""
-    tree = ast.parse(source_replace(source))
-    t1 = SyntaxTransformer()
-    t1.convert_xor = convert_xor
+def assigns_to_dict(assigns):
+    """赋值表达式转成字典"""
+    return {ast.unparse(a.targets): ast.unparse(a.value) for a in assigns}
+
+
+def raw_to_code(raw):
+    """导入语句转字符列表"""
+    return '\n'.join([ast.unparse(a) for a in raw])
+
+
+def sources_to_asts(*sources, convert_xor: bool):
+    """输入多份源代码"""
+
+    def _source_to_asts(source):
+        """源代码"""
+        tree = ast.parse(source_replace(source))
+
+        if isinstance(tree.body[0], ast.FunctionDef):
+            body = tree.body[0].body
+        else:
+            body = tree.body
+
+        return body
+
+    tree = ast.parse("")
+    for arg in sources:
+        tree.body.extend(_source_to_asts(arg))
+
+    t1 = SyntaxTransformer(convert_xor)
     t1.visit(tree)
-    t = RenameTransformer()
+    t = RenameTransformer({}, {})
     t.visit(tree)
 
     raw = []
     assigns = []
 
-    if isinstance(tree.body[0], ast.FunctionDef):
-        body = tree.body[0].body
-    else:
-        body = tree.body
-
-    for node in body:
+    for node in tree.body:
         # 特殊处理的节点
         if isinstance(node, ast.Assign):
             assigns.append(node)
@@ -345,16 +349,6 @@ def _source_to_asts(source, convert_xor: bool):
     return raw_to_code(raw), assigns_to_dict(assigns), t.funcs_new, t.args_new, t.targets_new
 
 
-def assigns_to_dict(assigns):
-    """赋值表达式转成字典"""
-    return {ast.unparse(a.targets): ast.unparse(a.value) for a in assigns}
-
-
-def raw_to_code(raw):
-    """导入语句转字符列表"""
-    return '\n'.join([ast.unparse(a) for a in raw])
-
-
 def _add_default_type(globals_):
     # 这种写法可以省去由用户导入Eq一类的工作
     globals_['Add'] = Add

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/expr_codegen/expr.py

@@ -116,7 +116,7 @@ def is_simple_expr(expr):
     return False
 
 
-def get_current_by_prefix(expr, **kwargs):
+def get_current_by_prefix(expr, date, asset, **kwargs):
     """表达式根节点信息。按名称前缀。例如
 
     OPEN取的是OPEN，得cl
@@ -130,16 +130,16 @@ def get_current_by_prefix(expr, **kwargs):
                 prefix2 = expr.name[:2]
 
                 if prefix2 == TS:
-                    return TS, 'asset'
+                    return TS, asset
                 if prefix2 == CS:
-                    return CS, 'date'
+                    return CS, date
                 if prefix2 == GP:
-                    return GP, 'date', expr.args[0].name
+                    return GP, date, expr.args[0].name
     # 不需分组
     return CL_TUP
 
 
-def get_current_by_name(expr, ts_names, cs_names, gp_names, **kwargs):
+def get_current_by_name(expr, ts_names, cs_names, gp_names, date, asset, **kwargs):
     """表达式根节点信息。按名称。
 
     Parameters
@@ -151,17 +151,21 @@ def get_current_by_name(expr, ts_names, cs_names, gp_names, **kwargs):
         横截面算子名称字符串集合
     gp_names
         分组算子名称字符串集合
+    date
+        日期字符串
+    asset
+        资产字符串
     kwargs
 
     """
     if expr.is_Function:
         if hasattr(expr, 'name'):  # Or 没有名字
             if expr.name in ts_names:
-                return TS, 'asset'
+                return TS, asset
             if expr.name in cs_names:
-                return CS, 'date'
+                return CS, date
             if expr.name in gp_names:
-                return GP, 'date', expr.args[0].name
+                return GP, date, expr.args[0].name
 
     # 不需分组
     return CL_TUP
@@ -171,7 +175,7 @@ def get_current_by_name(expr, ts_names, cs_names, gp_names, **kwargs):
 # __level__ = 0
 
 
-def get_children(func, func_kwargs, expr, output_exprs, output_symbols):
+def get_children(func, func_kwargs, expr, output_exprs, output_symbols, date, asset):
     """表达式整体信息。例如
 
     -ts_corr返回{ts}而不是 {cl}
@@ -190,6 +194,8 @@ def get_children(func, func_kwargs, expr, output_exprs, output_symbols):
         输出分割后的了表达式
     output_symbols
         输出每个子表达式中的符号
+    date
+    asset
 
     Returns
     -------
@@ -199,8 +205,8 @@ def get_children(func, func_kwargs, expr, output_exprs, output_symbols):
     # __level__ += 1
 
     try:
-        curr = func(expr, **func_kwargs)
-        children = [get_children(func, func_kwargs, a, output_exprs, output_symbols) for a in expr.args]
+        curr = func(expr, date, asset, **func_kwargs)
+        children = [get_children(func, func_kwargs, a, output_exprs, output_symbols, date, asset) for a in expr.args]
 
         # print(expr, curr, children, __level__)
         # if __level__ == 6:

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/expr_codegen/model.py

@@ -206,14 +206,14 @@ def create_dag_exprs(exprs):
     return G
 
 
-def init_dag_exprs(G, func, func_kwargs):
+def init_dag_exprs(G, func, func_kwargs, date, asset):
     """使用表达式信息初始化DAG"""
     for i, generation in enumerate(nx.topological_generations(G)):
         # print(i, generation)
         for node in generation:
             expr = G.nodes[node]['expr']
             syms = []
-            children = get_children(func, func_kwargs, expr, [], syms)
+            children = get_children(func, func_kwargs, expr, [], syms, date, asset)
             G.nodes[node]['children'] = children
             G.nodes[node]['key'] = get_key(children)
             G.nodes[node]['symbols'] = [str(s) for s in syms]
@@ -248,7 +248,11 @@ def merge_nodes_1(G: nx.DiGraph, keep_nodes, *args):
                     succ = G.succ[node]
                     # 下游只有一个，直接替换。
                     if len(succ) == 1:
-                        skip_expr_node(G, node, keep_nodes)
+                        for s in succ:
+                            # if_else(_A>_B,_A,_B)会出现量次，不能删
+                            if G.nodes[s]['symbols'].count(node) > 1:
+                                continue
+                            skip_expr_node(G, node, keep_nodes)
             else:
                 # 复制一次，防止修改后报错
                 for p in pred.copy():
@@ -263,7 +267,10 @@ def merge_nodes_1(G: nx.DiGraph, keep_nodes, *args):
                         succ = G.succ[p]
                         # 下游只有一个，直接替换。
                         if len(succ) == 1:
-                            skip_expr_node(G, p, keep_nodes)
+                            for s in succ:
+                                if G.nodes[s]['symbols'].count(p) > 1:
+                                    continue
+                                skip_expr_node(G, p, keep_nodes)
             next_pred.extend(pred)
         # 更新下一次循环
         this_pred = list(set(next_pred))
@@ -288,7 +295,10 @@ def merge_nodes_2(G: nx.DiGraph, keep_nodes, *args):
                 if len(succ) > 1:
                     # 上游节点只有一个下游，当前就是自己了
                     continue
-                skip_expr_node(G, p, keep_nodes)
+                for s in succ:
+                    if G.nodes[s]['symbols'].count(p) > 1:
+                        continue
+                    skip_expr_node(G, p, keep_nodes)
             # 只做根节点，所以没有下一次了
             # next_pred.extend(pred)
         # 更新下一次循环
@@ -352,16 +362,16 @@ def skip_expr_node(G: nx.DiGraph, node, keep_nodes):
     return G
 
 
-def dag_start(exprs_dict, func, func_kwargs):
+def dag_start(exprs_dict, func, func_kwargs, date, asset):
     """初始生成DAG"""
     G = create_dag_exprs(exprs_dict)
-    G = init_dag_exprs(G, func, func_kwargs)
+    G = init_dag_exprs(G, func, func_kwargs, date, asset)
 
     # 分层输出
     return G
 
 
-def dag_middle(G, exprs_names, func, func_kwargs):
+def dag_middle(G, exprs_names, func, func_kwargs, date, asset):
     """删除几个没有必要的节点"""
     G = remove_paths_by_zero_outdegree(G, exprs_names)
     # 以下划线开头的节点，不保留
@@ -370,7 +380,7 @@ def dag_middle(G, exprs_names, func, func_kwargs):
     G = merge_nodes_2(G, keep_nodes, *keep_nodes)
 
     # 由于表达式修改，需再次更新表达式
-    G = init_dag_exprs(G, func, func_kwargs)
+    G = init_dag_exprs(G, func, func_kwargs, date, asset)
 
     # 分层输出
     return G

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/expr_codegen/pandas/code.py

@@ -30,8 +30,8 @@ def get_groupby_from_tuple(tup, func_name, drop_cols):
 def symbols_to_code(syms, alias):
     a = [f"{s}" for s in syms]
     b = [f"'{alias.get(s, s)}'" for s in syms]
-    return f"""_ = ({','.join(b)},)
-({','.join(a)},) = _"""
+    return f"""_ = [{','.join(b)}]
+[{','.join(a)}] = _"""
 
 
 def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,

{expr_codegen-0.8.7/expr_codegen/polars → expr_codegen-0.9.1/expr_codegen/polars_group}/code.py

@@ -6,7 +6,7 @@ from jinja2 import FileSystemLoader, TemplateNotFound
 
 from expr_codegen.expr import TS, CS, GP
 from expr_codegen.model import ListDictList
-from expr_codegen.polars.printer import PolarsStrPrinter
+from expr_codegen.polars_group.printer import PolarsStrPrinter
 
 
 def get_groupby_from_tuple(tup, func_name, drop_cols):
@@ -31,8 +31,8 @@ def symbols_to_code(syms, alias):
     a = [f"{s}" for s in syms]
     b = [f"r'{alias.get(s, s)}'" for s in syms]  #
     b = [f"'{alias.get(s, s)}'" for s in syms]
-    return f"""_ = ({','.join(b)},)
-({','.join(a)},) = (pl.col(i) for i in _)"""
+    return f"""_ = [{','.join(b)}]
+[{','.join(a)}] = [pl.col(i) for i in _]"""
 
 
 def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,

{expr_codegen-0.8.7/expr_codegen/polars → expr_codegen-0.9.1/expr_codegen/polars_group}/printer.py

@@ -40,6 +40,11 @@ class PolarsStrPrinter(StrPrinter):
                                 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)
@@ -71,8 +76,8 @@ class PolarsStrPrinter(StrPrinter):
         PREC = PRECEDENCE["Mul"]
         return "~%s" % self.parenthesize(expr.args[0], PREC)
 
-    def _print_gp_rank(self, expr):
-        return "cs_rank(%s)" % self._print(expr.args[1])
-
-    def _print_gp_demean(self, expr):
-        return "cs_demean(%s)" % self._print(expr.args[1])
+    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))

expr_codegen-0.9.1/expr_codegen/polars_over/code.py

@@ -0,0 +1,113 @@
+import os
+from typing import Sequence, Dict
+
+import jinja2
+from jinja2 import FileSystemLoader, TemplateNotFound
+
+from expr_codegen.expr import TS, CS, GP
+from expr_codegen.model import ListDictList
+from expr_codegen.polars_over.printer import PolarsStrPrinter
+
+
+def get_groupby_from_tuple(tup, func_name, drop_cols):
+    """从传入的元组中生成分组运行代码"""
+    prefix2, *_ = tup
+
+    if prefix2 == TS:
+        # 组内需要按时间进行排序，需要维持顺序
+        prefix2, asset = tup
+        return f'df = {func_name}(df.sort(_ASSET_, _DATE_)).drop(*{drop_cols})'
+    if prefix2 == CS:
+        prefix2, date = tup
+        return f'df = {func_name}(df.sort(_DATE_)).drop(*{drop_cols})'
+    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).drop(*{drop_cols})'
+
+
+def symbols_to_code(syms, alias):
+    a = [f"{s}" for s in syms]
+    b = [f"r'{alias.get(s, s)}'" for s in syms]  #
+    b = [f"'{alias.get(s, 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='template.py.j2',
+            date='date', asset='asset',
+            alias: Dict[str, str] = {},
+            extra_codes: Sequence[str] = ()):
+    """基于模板的代码生成"""
+    # 打印Polars风格代码
+    p = PolarsStrPrinter()
+
+    # polars风格代码
+    funcs = {}
+    # 分组应用代码。这里利用了字典按插入顺序排序的特点，将排序放在最前
+    groupbys = {'sort': ''}
+    # 处理过后的表达式
+    exprs_dst = []
+    syms_out = []
+
+    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 = kv
+                    s1 = str(ex)
+                    s2 = p.doprint(ex)
+                    if s1 != s2:
+                        # 不想等，打印注释，显示会更直观察
+                        func_code.append(f"# {va} = {s1}")
+                    if k[0] == TS:
+                        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:
+                        func_code.append(f"{va}=({s2}).over(_DATE_, '{k[2]}'),")
+                    else:
+                        func_code.append(f"{va}={s2},")
+                    exprs_dst.append(f"{va} = {s1}")
+                    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, alias)
+    syms2 = symbols_to_code(syms_out, alias)
+
+    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.9.1/expr_codegen/polars_over/printer.py

@@ -0,0 +1,83 @@
+from sympy import Basic, Function, StrPrinter
+from sympy.printing.precedence import precedence, PRECEDENCE
+
+
+# TODO: 如有新添加函数，但表达式有变更才需要在此补充对应的打印代码，否则可以省略
+
+class PolarsStrPrinter(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):
+        return expr.name
+
+    def _print_Equality(self, expr):
+        PREC = precedence(expr)
+        return "%s==%s" % (self.parenthesize(expr.args[0], PREC), self.parenthesize(expr.args[1], PREC))
+
+    def _print_Or(self, expr):
+        PREC = PRECEDENCE["Mul"]
+        return " | ".join(self.parenthesize(arg, PREC) for arg in expr.args)
+
+    def _print_Xor(self, expr):
+        PREC = PRECEDENCE["Mul"]
+        return " ^ ".join(self.parenthesize(arg, PREC) for arg in expr.args)
+
+    def _print_And(self, expr):
+        PREC = PRECEDENCE["Mul"]
+        return " & ".join(self.parenthesize(arg, PREC) for arg in expr.args)
+
+    def _print_Not(self, expr):
+        PREC = PRECEDENCE["Mul"]
+        return "~%s" % self.parenthesize(expr.args[0], PREC)
+
+    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))

expr_codegen-0.9.1/expr_codegen/polars_over/template.py.j2

@@ -0,0 +1,79 @@
+# this code is auto generated by the expr_codegen
+# https://github.com/wukan1986/expr_codegen
+# 此段代码由 expr_codegen 自动生成，欢迎提交 issue 或 pull request
+
+import numpy as np  # noqa
+import pandas as pd  # noqa
+import polars as pl  # noqa
+import polars.selectors as cs  # noqa
+from loguru import logger  # noqa
+
+# ===================================
+# 导入优先级，例如：ts_RSI在ta与talib中都出现了，优先使用ta
+# 运行时，后导入覆盖前导入，但IDE智能提示是显示先导入的
+_ = 0  # 只要之前出现了语句，之后的import位置不参与调整
+# from polars_ta.prefix.talib import *  # noqa
+from polars_ta.prefix.tdx import *  # noqa
+from polars_ta.prefix.ta import *  # noqa
+from polars_ta.prefix.wq import *  # noqa
+from polars_ta.prefix.cdl import *  # noqa
+
+# ===================================
+
+{{ syms1 }}
+
+{{ syms2 }}
+
+_DATE_ = '{{ date }}'
+_ASSET_ = '{{ asset }}'
+_NONE_ = None
+_TRUE_ = True
+_FALSE_ = False
+
+{%-for row in extra_codes %}
+{{ row-}}
+{% endfor %}
+
+{% for key, value in funcs.items() %}
+def {{ key }}(df: pl.DataFrame) -> pl.DataFrame:
+{{ value }}
+    return df
+{% endfor %}
+
+"""
+{%-for row in exprs_dst %}
+{{ row-}}
+{% endfor %}
+"""
+
+"""
+{%-for key, value in exprs_src.items() %}
+{{ key }} = {{ value-}}
+{% endfor %}
+"""
+
+
+def main(df: pl.DataFrame) -> pl.DataFrame:
+    # logger.info("start...")
+    {% for key, value in groupbys.items() %}
+    {{ value-}}
+    {% endfor %}
+
+    # drop intermediate columns
+    # df = df.select(pl.exclude(r'^_x_\d+$'))
+    df = df.select(~cs.starts_with("_"))
+
+    # shrink
+    df = df.select(cs.all().shrink_dtype())
+    df = df.shrink_to_fit()
+
+    # logger.info('done')
+
+    # save
+    # df.write_parquet('output.parquet')
+
+    return df
+
+if __name__ in ("__main__", "builtins"):
+    # TODO: 数据加载或外部传入
+    df_output = main(df_input)

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/expr_codegen/tool.py

@@ -53,7 +53,7 @@ class ExprTool:
         self.get_current_func = func
         self.get_current_func_kwargs = kwargs
 
-    def extract(self, expr):
+    def extract(self, expr, date, asset):
         """抽取分割后的子公式
 
         Parameters
@@ -73,12 +73,13 @@ class ExprTool:
         syms = []
         get_children(self.get_current_func, self.get_current_func_kwargs,
                      expr,
-                     output_exprs=exprs, output_symbols=syms)
+                     output_exprs=exprs, output_symbols=syms,
+                     date=date, asset=asset)
         # print('=' * 20, expr)
         # print(exprs)
         return exprs, syms
 
-    def merge(self, **kwargs):
+    def merge(self, date, asset, **kwargs):
         """合并多个表达式
 
         1. 先抽取分割子公式
@@ -93,7 +94,7 @@ class ExprTool:
         -------
         表达式列表
         """
-        exprs_syms = [self.extract(v) for v in kwargs.values()]
+        exprs_syms = [self.extract(v, date, asset) for v in kwargs.values()]
         exprs = []
         syms = []
         for e, s in exprs_syms:
@@ -164,14 +165,14 @@ class ExprTool:
 
         return self.exprs_dict
 
-    def dag(self, merge: bool):
+    def dag(self, merge: bool, date, asset):
         """生成DAG"""
-        G = dag_start(self.exprs_dict, self.get_current_func, self.get_current_func_kwargs)
+        G = dag_start(self.exprs_dict, 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)
+            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: str = 'polars', template_file: str = 'template.py.j2',
+    def all(self, exprs_src, style: str = '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] = {},
@@ -183,7 +184,7 @@ class ExprTool:
         exprs_src: dict
             表达式字典
         style: str
-            代码风格。可选值 ('polars', 'pandas')
+            代码风格。可选值 ('polars_group', 'polars_over', 'pandas')
         template_file: str
             根据需求可定制模板
         replace:bool
@@ -206,24 +207,26 @@ class ExprTool:
         代码字符串
 
         """
-        assert style in ('polars', 'pandas')
+        assert style in ('polars_group', 'polars_over', 'pandas')
 
         if replace:
             exprs_src = replace_exprs(exprs_src)
 
         # 子表达式在前，原表式在最后
-        exprs_dst, syms_dst = self.merge(**exprs_src)
+        exprs_dst, syms_dst = self.merge(date, asset, **exprs_src)
 
         # 提取公共表达式
         self.cse(exprs_dst, symbols_repl=numbered_symbols('_x_'), symbols_redu=exprs_src.keys())
         # 有向无环图流转
-        exprs_ldl, G = self.dag(True)
+        exprs_ldl, G = self.dag(True, date, asset)
 
         if regroup:
             exprs_ldl.optimize()
 
-        if style == 'polars':
-            from expr_codegen.polars.code import codegen
+        if style == 'polars_group':
+            from expr_codegen.polars_group.code import codegen
+        elif style == 'polars_over':
+            from expr_codegen.polars_over.code import codegen
         else:
             from expr_codegen.pandas.code import codegen
 
@@ -257,7 +260,7 @@ class ExprTool:
                   source: str, *more_sources: str,
                   extra_codes: str, output_file: str,
                   convert_xor: bool,
-                  style='polars', template_file='template.py.j2',
+                  style='polars_over', template_file='template.py.j2',
                   date='date', asset='asset') -> str:
         """通过字符串生成代码， 加了缓存，多次调用不重复生成"""
         raw, exprs_dict = sources_to_exprs(self.globals_, source, *more_sources, convert_xor=convert_xor)
@@ -288,7 +291,7 @@ def codegen_exec(df,
                  extra_codes: str = r'CS_SW_L1 = pl.col(r"^sw_l1_\d+$")',
                  output_file: Optional[str] = None,
                  convert_xor: bool = False,
-                 style: str = 'polars', template_file: str = 'template.py.j2',
+                 style: str = 'polars_over', template_file: str = 'template.py.j2',
                  date: str = 'date', asset: str = 'asset'
                  ):
     """快速转换源代码并执行
@@ -306,7 +309,7 @@ def codegen_exec(df,
     convert_xor: bool
         ^ 转成异或还是乘方
     style: str
-        代码风格。可选值 ('polars', 'pandas')
+        代码风格。可选值 ('polars_group', 'polars_over', 'pandas')
     template_file: str
         代码模板
     date: str

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/expr_codegen.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: expr_codegen
-Version: 0.8.7
+Version: 0.9.1
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License
@@ -206,6 +206,7 @@ df = codegen_exec(df, _code_block_1, _code_block_2)  # 只执行，不保存代
 6. 不支持`A<=B<=C`，需手工替换成`(A<=B)&(B<=C)`
 7. 支持`A[0]+B[1]+C[2]`，底层会转成`A+ts_delay(B,1)+ts_delay(C,2)`
 8. 支持`~A`,底层会转换成`Not(A)`
+9. `gp_`开头的函数都会返回对应的`cs_`函数。如`gp_func(A,B,C)`会替换成`cs_func(B,C)`,其中`A`用在了`groupby([date, A])`
 
 ## 下划线开头的变量
 

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/expr_codegen.egg-info/SOURCES.txt

@@ -19,7 +19,11 @@ expr_codegen/pandas/__init__.py
 expr_codegen/pandas/code.py
 expr_codegen/pandas/printer.py
 expr_codegen/pandas/template.py.j2
-expr_codegen/polars/__init__.py
-expr_codegen/polars/code.py
-expr_codegen/polars/printer.py
-expr_codegen/polars/template.py.j2
+expr_codegen/polars_group/__init__.py
+expr_codegen/polars_group/code.py
+expr_codegen/polars_group/printer.py
+expr_codegen/polars_group/template.py.j2
+expr_codegen/polars_over/__init__.py
+expr_codegen/polars_over/code.py
+expr_codegen/polars_over/printer.py
+expr_codegen/polars_over/template.py.j2

{expr_codegen-0.8.7 → expr_codegen-0.9.1}/pyproject.toml

@@ -38,12 +38,14 @@ packages = [
     "expr_codegen",
     "expr_codegen.latex",
     "expr_codegen.pandas",
-    "expr_codegen.polars",
+    "expr_codegen.polars_group",
+    "expr_codegen.polars_over",
 ]
 
 [tool.setuptools.package-data]
 "expr_codegen.pandas" = ["*.j2"]
-"expr_codegen.polars" = ["*.j2"]
+"expr_codegen.polars_group" = ["*.j2"]
+"expr_codegen.polars_over" = ["*.j2"]
 
 [tool.setuptools.dynamic]
 version = { attr = "expr_codegen._version.__version__" }

expr_codegen-0.8.7/expr_codegen/_version.py

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