Coverage for graphica/core/safe_eval.py: 93.3%

1"""利用者が入力した式を eval / exec を使わずに評価する(AST を辿り、許した構文だけ通す)。 

2 

3使えるのは数値・変数・四則演算・べき乗・剰余・比較・and / or / not(要素ごとの & / | / ~)と許した関数。 

4列の計算では ALLOWED_SERIES_METHODS のメソッド(A.mean()、A.rolling(5).mean() など)も使える。 

5属性の任意呼び出し・添字・lambda・内包表記・import・_ で始まる名前は拒否する。 

6列の計算では、空白や記号を含む列名を `強度 (a.u.)` のようにバッククォートで囲んで書ける(pandas の query/eval と同じ)。 

7""" 

8import ast 

9import keyword 

10import operator 

11import re 

12 

13import numpy as np 

14from typing import Any, Callable 

15 

16 

17class SafeEvalError(ValueError): 

18 """式に許していない構文・名前・関数・属性がある。""" 

19 

20 

21DEFAULT_FUNCTIONS: dict[str, Any] = { 

22 'exp': np.exp, 'log': np.log, 'log10': np.log10, 'sqrt': np.sqrt, 

23 'sin': np.sin, 'cos': np.cos, 'tan': np.tan, 'abs': np.abs, 

24 'pi': np.pi, 'e': np.e, 

25} 

26 

27# 列の計算でだけ許すメソッド。I/O、任意のコールバックを取るもの(apply など)、ダンダー名は入れない 

28ALLOWED_SERIES_METHODS = frozenset({ 

29 'mean', 'std', 'var', 'min', 'max', 'median', 'sum', 'abs', 

30 'diff', 'cumsum', 'cumprod', 'cummax', 'cummin', 'pct_change', 

31 'shift', 'rolling', 'expanding', 'ewm', 'rank', 'round', 'clip', 

32 'quantile', 'skew', 'kurt', 'count', 'size', 

33}) 

34 

35_BINOPS: dict[type, Callable[..., Any]] = { 

36 ast.Add: operator.add, ast.Sub: operator.sub, ast.Mult: operator.mul, 

37 ast.Div: operator.truediv, ast.Pow: operator.pow, ast.Mod: operator.mod, 

38 ast.FloorDiv: operator.floordiv, 

39} 

40# Series に素の not / and / or を使うと「真偽値が曖昧」になるので、pandas.eval と同じく ~ / & / | として扱う 

41_UNARYOPS: dict[type, Callable[..., Any]] = {ast.UAdd: operator.pos, ast.USub: operator.neg, ast.Not: operator.invert} 

42_BOOLOPS: dict[type, Callable[[Any, Any], Any]] = {ast.And: operator.and_, ast.Or: operator.or_} 

43_COMPAREOPS: dict[type, Callable[[Any, Any], Any]] = { 

44 ast.Eq: operator.eq, ast.NotEq: operator.ne, 

45 ast.Lt: operator.lt, ast.LtE: operator.le, 

46 ast.Gt: operator.gt, ast.GtE: operator.ge, 

47} 

48 

49 

50def _parse(formula: str) -> ast.expr: 

51 try: 

52 tree = ast.parse(formula, mode='eval') 

53 except SyntaxError as e: 

54 raise SafeEvalError(f"数式の構文が不正です: {e}") from e 

55 return tree.body 

56 

57 

58class _Evaluator: 

59 def __init__(self, variables: dict[str, Any], functions: dict[str, Any], 

60 allowed_methods: frozenset[str] | None) -> None: 

61 self.variables = variables 

62 self.functions = functions 

63 # None なら属性アクセスを一切許さない(フィットの式) 

64 self.allowed_methods = allowed_methods 

65 

66 def eval(self, node: ast.AST) -> Any: 

67 if isinstance(node, ast.Constant): 

68 if isinstance(node.value, (int, float, complex)): 

69 return node.value 

70 raise SafeEvalError(f"許可されていない定数です: {node.value!r}") 

71 

72 if isinstance(node, ast.Name): 

73 if node.id in self.variables: 

74 return self.variables[node.id] 

75 if node.id in self.functions: 

76 return self.functions[node.id] 

77 raise SafeEvalError(f"未定義の変数/関数です: {node.id}") 

78 

79 if isinstance(node, ast.BinOp): 

80 op = _BINOPS.get(type(node.op)) 

81 if op is None: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 raise SafeEvalError(f"許可されていない演算子です: {type(node.op).__name__}") 

83 return op(self.eval(node.left), self.eval(node.right)) 

84 

85 if isinstance(node, ast.UnaryOp): 

86 op = _UNARYOPS.get(type(node.op)) 

87 if op is None: 87 ↛ 88line 87 didn't jump to line 88 because the condition on line 87 was never true

88 raise SafeEvalError(f"許可されていない単項演算子です: {type(node.op).__name__}") 

89 return op(self.eval(node.operand)) 

90 

91 if isinstance(node, ast.Compare): 

92 left = self.eval(node.left) 

93 result = None 

94 for op_node, comparator in zip(node.ops, node.comparators): 

95 op = _COMPAREOPS.get(type(op_node)) 

96 if op is None: 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true

97 raise SafeEvalError(f"許可されていない比較演算子です: {type(op_node).__name__}") 

98 right = self.eval(comparator) 

99 partial = op(left, right) 

100 result = partial if result is None else (result & partial) 

101 left = right 

102 return result 

103 

104 if isinstance(node, ast.BoolOp): 

105 op = _BOOLOPS.get(type(node.op)) 

106 if op is None: 106 ↛ 107line 106 didn't jump to line 107 because the condition on line 106 was never true

107 raise SafeEvalError(f"許可されていない論理演算子です: {type(node.op).__name__}") 

108 values = [self.eval(v) for v in node.values] 

109 result = values[0] 

110 for v in values[1:]: 

111 result = op(result, v) 

112 return result 

113 

114 if isinstance(node, ast.Attribute): 

115 if self.allowed_methods is None: 

116 raise SafeEvalError("この数式では属性アクセスは使用できません。") 

117 if node.attr.startswith('_') or node.attr not in self.allowed_methods: 

118 raise SafeEvalError(f"許可されていないメソッド/属性です: {node.attr}") 

119 return getattr(self.eval(node.value), node.attr) 

120 

121 if isinstance(node, ast.Call): 

122 args = [self.eval(a) for a in node.args] 

123 # **expr の展開は kw.arg が None。拒否しないと関数呼び出しで TypeError が漏れる 

124 for kw in node.keywords: 

125 if kw.arg is None: 125 ↛ 124line 125 didn't jump to line 124 because the condition on line 125 was always true

126 raise SafeEvalError("この形式の関数呼び出し(**による引数展開)は許可されていません。") 

127 kwargs = {str(kw.arg): self.eval(kw.value) for kw in node.keywords} 

128 if isinstance(node.func, ast.Name): 

129 name = node.func.id 

130 if name not in self.functions or not callable(self.functions[name]): 

131 raise SafeEvalError(f"許可されていない関数です: {name}") 

132 return self.functions[name](*args, **kwargs) 

133 if isinstance(node.func, ast.Attribute): 

134 method = self.eval(node.func) 

135 if not callable(method): 135 ↛ 136line 135 didn't jump to line 136 because the condition on line 135 was never true

136 raise SafeEvalError(f"'{node.func.attr}' は呼び出せません。") 

137 return method(*args, **kwargs) 

138 raise SafeEvalError("この形式の関数呼び出しは許可されていません。") 

139 

140 raise SafeEvalError(f"許可されていない構文です: {type(node).__name__}") 

141 

142 

143def safe_eval_formula(formula: str, variables: dict[str, Any], functions: dict[str, Any] | None = None) -> Any: 

144 """フィットの式を評価する。属性アクセスとメソッド呼び出しは許さない。functions は DEFAULT_FUNCTIONS に足す。""" 

145 funcs = dict(DEFAULT_FUNCTIONS) 

146 if functions: 146 ↛ 147line 146 didn't jump to line 147 because the condition on line 146 was never true

147 funcs.update(functions) 

148 node = _parse(formula) 

149 return _Evaluator(variables, funcs, allowed_methods=None).eval(node) 

150 

151 

152_BACKTICK_COLUMN = re.compile(r"`([^`]*)`") 

153 

154 

155def column_reference(name: str) -> str: 

156 """列の計算の式でその列を指す書き方。名前のまま書けない列名はバッククォートで囲む。""" 

157 if name.isidentifier() and not keyword.iskeyword(name) and not name.startswith('_'): 

158 return name 

159 return f"`{name}`" 

160 

161 

162def _replace_backtick_columns(formula: str, variables: dict[str, Any]) -> str: 

163 """`列名` を、ほかの名前とぶつからない変数名に置き換え、その変数を variables に足す。""" 

164 taken = set(variables) | set(DEFAULT_FUNCTIONS) 

165 placeholders: dict[str, str] = {} 

166 

167 def replace(match: re.Match[str]) -> str: 

168 name = match.group(1) 

169 if name not in placeholders: 

170 if name not in variables: 

171 raise SafeEvalError(f"未定義の列です: {name}") 

172 index = len(placeholders) 

173 while f"backtick_column_{index}" in taken: 

174 index += 1 

175 placeholder = f"backtick_column_{index}" 

176 taken.add(placeholder) 

177 placeholders[name] = placeholder 

178 return placeholders[name] 

179 

180 replaced = _BACKTICK_COLUMN.sub(replace, formula) 

181 if "`" in replaced: 

182 raise SafeEvalError("列名を囲むバッククォート ` が閉じていません。") 

183 for name, placeholder in placeholders.items(): 

184 variables[placeholder] = variables[name] 

185 return replaced 

186 

187 

188def safe_eval_column_formula(df: Any, formula: str) -> Any: 

189 """列名を変数として列の計算式を評価する(ALLOWED_SERIES_METHODS のメソッドも使える)。""" 

190 variables = {str(col): df[col] for col in df.columns} 

191 if "`" in formula: 

192 formula = _replace_backtick_columns(formula, variables) 

193 node = _parse(formula) 

194 return _Evaluator(variables, dict(DEFAULT_FUNCTIONS), allowed_methods=ALLOWED_SERIES_METHODS).eval(node)