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

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

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

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

{expr_codegen-0.5.2 → expr_codegen-0.6.0}/expr_codegen/codes.py

@@ -1,6 +1,8 @@
 import ast
 import re
 
+from sympy import Add, Mul, Pow, Eq
+
 from expr_codegen.expr import register_symbols, dict_to_exprs
 
 
@@ -191,8 +193,20 @@ def raw_to_code(raw):
     return '\n'.join([ast.unparse(a) for a in raw])
 
 
+def _add_default_type(globals_):
+    # 这种写法可以省去由用户导入Eq一类的工作
+    globals_['Add'] = Add
+    globals_['Mul'] = Mul
+    globals_['Pow'] = Pow
+    globals_['Eq'] = Eq
+    return globals_
+
+
 def sources_to_exprs(globals_, *sources, safe: bool = True):
     """将源代码转换成表达式"""
+
+    globals_ = _add_default_type(globals_)
+
     raw, assigns, funcs_new, args_new, targets_new = sources_to_asts(*sources)
     register_symbols(funcs_new, globals_, is_function=True)
     register_symbols(args_new, globals_, is_function=False)

{expr_codegen-0.5.2 → expr_codegen-0.6.0}/expr_codegen/pandas/code.py

@@ -1,5 +1,5 @@
 import os
-from typing import Sequence
+from typing import Sequence, Dict
 
 import jinja2
 from jinja2 import FileSystemLoader, TemplateNotFound
@@ -27,9 +27,9 @@ def get_groupby_from_tuple(tup, func_name):
     return f'df = {func_name}(df)'
 
 
-def symbols_to_code(syms):
+def symbols_to_code(syms, alias):
     a = [f"{s}" for s in syms]
-    b = [f"'{s}'" for s in syms]
+    b = [f"r'{alias.get(s, s)}'" for s in syms]
     return f"""_ = ({','.join(b)},)
 ({','.join(a)},) = _"""
 
@@ -37,6 +37,7 @@ def symbols_to_code(syms):
 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] = ()):
     """基于模板的代码生成"""
     # 打印Pandas风格代码
@@ -65,7 +66,8 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
                     va, ex = kv
                     func_code.append(f"    # {va} = {ex}\n    df[{va}] = {p.doprint(ex)}")
                     exprs_dst.append(f"{va} = {ex}")
-                    syms_out.append(va)
+                    if va not in syms_dst:
+                        syms_out.append(va)
 
             if k[0] == TS:
                 groupbys['sort'] = f'df = df.sort_values(by=[_DATE_, _ASSET_]).reset_index(drop=True)'
@@ -77,8 +79,8 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
             # 分组应用代码
             groupbys[func_name] = get_groupby_from_tuple(k, func_name)
 
-    syms1 = symbols_to_code(syms_dst)
-    syms2 = symbols_to_code(syms_out)
+    syms1 = symbols_to_code(syms_dst, alias)
+    syms2 = symbols_to_code(syms_out, alias)
 
     try:
         env = jinja2.Environment(loader=FileSystemLoader(os.path.dirname(__file__)))

{expr_codegen-0.5.2 → expr_codegen-0.6.0}/expr_codegen/pandas/template.py.j2

@@ -49,7 +49,7 @@ def {{ key }}(df: pd.DataFrame) -> pd.DataFrame:
 """
 
 
-def main(df: pd.DataFrame):
+def main(df: pd.DataFrame) -> pd.DataFrame:
     # logger.info("start...")
     {% for key, value in groupbys.items() %}
     {{ value-}}

{expr_codegen-0.5.2 → expr_codegen-0.6.0}/expr_codegen/polars/code.py

@@ -1,5 +1,5 @@
 import os
-from typing import Sequence
+from typing import Sequence, Dict
 
 import jinja2
 from jinja2 import FileSystemLoader, TemplateNotFound
@@ -27,9 +27,9 @@ def get_groupby_from_tuple(tup, func_name):
     return f'df = {func_name}(df)'
 
 
-def symbols_to_code(syms):
+def symbols_to_code(syms, alias):
     a = [f"{s}" for s in syms]
-    b = [f"'{s}'" for s in syms]
+    b = [f"r'{alias.get(s, s)}'" for s in syms]
     return f"""_ = ({','.join(b)},)
 ({','.join(a)},) = (pl.col(i) for i in _)"""
 
@@ -37,6 +37,7 @@ def symbols_to_code(syms):
 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风格代码
@@ -73,7 +74,8 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
 
                     func_code.append(f"{va}={s2},")
                     exprs_dst.append(f"{va} = {s1}")
-                    syms_out.append(va)
+                    if va not in syms_dst:
+                        syms_out.append(va)
             func_code.append(f"    )")
             func_code = func_code[1:]
 
@@ -87,8 +89,8 @@ def codegen(exprs_ldl: ListDictList, exprs_src, syms_dst,
             # 分组应用代码
             groupbys[func_name] = get_groupby_from_tuple(k, func_name)
 
-    syms1 = symbols_to_code(syms_dst)
-    syms2 = symbols_to_code(syms_out)
+    syms1 = symbols_to_code(syms_dst, alias)
+    syms2 = symbols_to_code(syms_out, alias)
 
     try:
         env = jinja2.Environment(loader=FileSystemLoader(os.path.dirname(__file__)))

{expr_codegen-0.5.2 → expr_codegen-0.6.0}/expr_codegen/polars/template.py.j2

@@ -51,7 +51,7 @@ def {{ key }}(df: pl.DataFrame) -> pl.DataFrame:
 """
 
 
-def main(df: pl.DataFrame):
+def main(df: pl.DataFrame) -> pl.DataFrame:
     # logger.info("start...")
     {% for key, value in groupbys.items() %}
     {{ value-}}

{expr_codegen-0.5.2 → expr_codegen-0.6.0}/expr_codegen/tool.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Sequence
+from typing import Sequence, Dict
 
 from black import Mode, format_str
 from sympy import simplify, cse, symbols, numbered_symbols
@@ -73,7 +73,7 @@ class ExprTool:
         exprs = exprs + list(kwargs.values())
 
         # print(exprs)
-
+        syms = [str(s) for s in syms]
         return exprs, syms
 
     def reduce(self, repl, redu):
@@ -141,6 +141,7 @@ class ExprTool:
     def all(self, exprs_src, style: str = 'polars', template_file: str = 'template.py.j2',
             replace: bool = True, regroup: bool = False, format: bool = True,
             date='date', asset='asset',
+            alias: Dict[str, str] = {},
             extra_codes: Sequence[object] = ()):
         """功能集成版，将几个功能写到一起方便使用
 
@@ -162,6 +163,8 @@ class ExprTool:
             日期字段名
         asset:str
             资产字段名
+        alias: Dict[str,str]
+            符号别名。可以变通的传入正则符号名
         extra_codes: Sequence[object]
             需要复制到模板中的额外代码
 
@@ -195,6 +198,7 @@ class ExprTool:
 
         codes = codegen(exprs_ldl, exprs_src, syms_dst,
                         filename=template_file, date=date, asset=asset,
+                        alias=alias,
                         extra_codes=extra_codes)
 
         if format:

{expr_codegen-0.5.2 → expr_codegen-0.6.0}/expr_codegen.egg-info/PKG-INFO

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

expr_codegen-0.5.2/expr_codegen/_version.py

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