expr-codegen 0.6.1__tar.gz → 0.6.3__tar.gz

{expr_codegen-0.6.1 → expr_codegen-0.6.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: expr_codegen
-Version: 0.6.1
+Version: 0.6.3
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License

expr_codegen-0.6.3/expr_codegen/_version.py

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

{expr_codegen-0.6.1 → expr_codegen-0.6.3}/expr_codegen/codes.py

@@ -1,5 +1,6 @@
 import ast
 import re
+from ast import expr
 
 from sympy import Add, Mul, Pow, Eq
 
@@ -21,7 +22,7 @@ class SympyTransformer(ast.NodeTransformer):
     # 映射
     funcs_map = {}
     args_map = {}
-    targets_map = {}
+    targets_map = {}  # 只对非下划线开头的生效
 
     def config_map(self, funcs_map, args_map, targets_map):
         self.funcs_map = funcs_map
@@ -43,18 +44,33 @@ class SympyTransformer(ast.NodeTransformer):
         self.generic_visit(node)
         return node
 
+    def __visit_Assign(self, target: expr):
+        old_target_id = target.id
+        new_target_id = self.targets_map.get(old_target_id, old_target_id)
+        self.targets_old.add(old_target_id)
+
+        # 赋值给下划线开头代码时，对其进行重命名，方便重复书写表达式时不冲突
+        if old_target_id.startswith('_'):
+            new_target_id = f'{old_target_id}_{len(self.targets_new):03d}'
+
+        if old_target_id != new_target_id:
+            self.targets_new.add(new_target_id)
+            target.id = new_target_id
+            # 记录修改的变量名，之后会使用到
+            self.args_map[old_target_id] = new_target_id
+
     def visit_Assign(self, node):
+        # 调整位置，支持循环赋值
+        # _A = _A+1 调整成 _A_001 = _A_000 + 1
+        self.generic_visit(node)
+
         # 提取输出变量名
         for target in node.targets:
             if isinstance(target, ast.Tuple):
                 for t in target.elts:
-                    self.targets_old.add(t.id)
-                    t.id = self.targets_map.get(t.id, t.id)
-                    self.targets_new.add(t.id)
+                    self.__visit_Assign(t)
             else:
-                self.targets_old.add(target.id)
-                target.id = self.targets_map.get(target.id, target.id)
-                self.targets_new.add(target.id)
+                self.__visit_Assign(target)
 
         # 处理 alpha=close 这种情况
         if isinstance(node.value, ast.Name):
@@ -62,7 +78,6 @@ class SympyTransformer(ast.NodeTransformer):
             node.value.id = self.args_map.get(node.value.id, node.value.id)
             self.args_new.add(node.value.id)
 
-        self.generic_visit(node)
         return node
 
     def visit_Compare(self, node):

{expr_codegen-0.6.1 → expr_codegen-0.6.3}/expr_codegen/pandas/template.py.j2

@@ -56,7 +56,8 @@ def main(df: pd.DataFrame) -> pd.DataFrame:
     {% endfor %}
 
     # drop intermediate columns
-    df = df.drop(columns=list(filter(lambda x: re.search(r"^_x_\d+", x), df.columns)))
+    # df = df.drop(columns=list(filter(lambda x: re.search(r"^_x_\d+", x), df.columns)))
+    df = df.drop(columns=list(filter(lambda x: x.startswith("_"), df.columns)))
 
     # logger.info('done')
 

{expr_codegen-0.6.1 → expr_codegen-0.6.3}/expr_codegen/polars/template.py.j2

@@ -58,7 +58,8 @@ def main(df: pl.DataFrame) -> pl.DataFrame:
     {% endfor %}
 
     # drop intermediate columns
-    df = df.select(pl.exclude(r'^_x_\d+$'))
+    # df = df.select(pl.exclude(r'^_x_\d+$'))
+    df = df.select(~cs.starts_with("_"))
 
     # shrink
     df = df.select(cs.all().shrink_dtype())

{expr_codegen-0.6.1 → expr_codegen-0.6.3}/expr_codegen/tool.py

@@ -1,9 +1,11 @@
 import inspect
-from typing import Sequence, Dict
+from functools import lru_cache
+from typing import Sequence, Dict, Optional
 
 from black import Mode, format_str
 from sympy import simplify, cse, symbols, numbered_symbols
 
+from expr_codegen.codes import sources_to_exprs
 from expr_codegen.expr import get_current_by_prefix, get_children, replace_exprs
 from expr_codegen.model import dag_start, dag_end, dag_middle
 
@@ -15,6 +17,7 @@ class ExprTool:
         self.get_current_func_kwargs = {}
         self.exprs_dict = {}
         self.exprs_names = []
+        self.globals_ = {}
 
     def set_current(self, func, **kwargs):
         self.get_current_func = func
@@ -206,3 +209,55 @@ class ExprTool:
             codes = format_str(codes, mode=Mode(line_length=600, magic_trailing_comma=True))
 
         return codes, G
+
+    def exec(self, codes: str, df_input):
+        """执行代码
+
+        Notes
+        -----
+        注意生成的代码已经约定输入用df_input，输出用df_output
+
+        """
+        globals_ = {'df_input': df_input}
+        exec(codes, globals_)
+        return globals_['df_output']
+
+    @lru_cache(maxsize=64)
+    def _get_codes(self, source: str, extra_codes: str, output_file: str) -> str:
+        """通过字符串生成代码， 加了缓存，多次调用不重复生成"""
+        raw, exprs_dict = sources_to_exprs(self.globals_, source, safe=False)
+
+        # 生成代码
+        codes, G = _TOOL_.all(exprs_dict, style='polars', template_file='template.py.j2',
+                              replace=True, regroup=True, format=True,
+                              date='date', asset='asset',
+                              # 复制了需要使用的函数，还复制了最原始的表达式
+                              extra_codes=(raw,
+                                           # 传入多个列的方法
+                                           extra_codes,
+                                           ))
+
+        if output_file is not None:
+            with open(output_file, 'w', encoding='utf-8') as f:
+                f.write(codes)
+
+        return codes
+
+
+_TOOL_ = ExprTool()
+
+
+def codegen_exec(globals_, code_block, df_input,
+                 extra_codes: str = r'CS_SW_L1 = pl.col(r"^sw_l1_\d+$")',
+                 output_file: Optional[str] = None):
+    """快速转换源代码并执行"""
+    _TOOL_.globals_ = globals_
+
+    if isinstance(code_block, str):
+        source = code_block
+    else:
+        source = inspect.getsource(code_block)
+
+    codes = _TOOL_._get_codes(source, extra_codes, output_file)
+
+    return _TOOL_.exec(codes, df_input)

{expr_codegen-0.6.1 → expr_codegen-0.6.3}/expr_codegen.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: expr_codegen
-Version: 0.6.1
+Version: 0.6.3
 Summary: symbol expression to polars expression tool
 Author-email: wukan <wu-kan@163.com>
 License: BSD 3-Clause License

expr_codegen-0.6.1/expr_codegen/_version.py

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