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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: expr_codegen
-Version: 0.6.0
+Version: 0.6.2
 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

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

{expr_codegen-0.6.0 → expr_codegen-0.6.2}/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,35 @@ 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):
         # 提取输出变量名
         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):
+            self.args_old.add(node.value.id)
+            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

{expr_codegen-0.6.0 → expr_codegen-0.6.2}/expr_codegen/expr.py

@@ -352,7 +352,7 @@ def _replace__repeat(e):
         node_name = get_node_name(node)
         node_args0_name = get_node_name(node.args[0])
         if node_name == node_args0_name:
-            if node.name in ('cs_rank', 'sign', 'Abs'):
+            if node.name in ('cs_rank', 'sign', 'Abs', 'abs_'):
                 replacements.append((node, node.args[0]))
     for node, replacement in replacements:
         print(node, '  ->  ', replacement)

{expr_codegen-0.6.0 → expr_codegen-0.6.2}/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.0 → expr_codegen-0.6.2}/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.0 → expr_codegen-0.6.2}/expr_codegen.egg-info/PKG-INFO

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

expr_codegen-0.6.0/expr_codegen/_version.py

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