expr-codegen 0.6.2__tar.gz → 0.6.4__tar.gz

{expr_codegen-0.6.2 → expr_codegen-0.6.4}/PKG-INFO

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

expr_codegen-0.6.4/expr_codegen/_version.py

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

{expr_codegen-0.6.2 → expr_codegen-0.6.4}/expr_codegen/codes.py

@@ -60,6 +60,10 @@ class SympyTransformer(ast.NodeTransformer):
             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):
@@ -74,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.2 → expr_codegen-0.6.4}/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,58 @@ 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(code_block, df_input,
+                 extra_codes: str = r'CS_SW_L1 = pl.col(r"^sw_l1_\d+$")',
+                 output_file: Optional[str] = None):
+    """快速转换源代码并执行"""
+    # 此代码来自于sympy.var
+    frame = inspect.currentframe().f_back
+    _TOOL_.globals_ = frame.f_globals.copy()
+    del frame
+
+    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.2 → expr_codegen-0.6.4}/expr_codegen.egg-info/PKG-INFO

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

expr_codegen-0.6.2/expr_codegen/_version.py

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