1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
|
from ast import Module, walk, Name, BinOp, Constant, operator, Mult, Div, Add, Sub, Pow, Assign
from types import CodeType
from typing import FrozenSet, Optional, cast
from .mod import Mod
class CodeOp(object):
result: str
parameters: FrozenSet[str]
variables: FrozenSet[str]
code: Module
operator: Optional[operator]
compiled: CodeType
def __init__(self, code: Module):
self.code = code
assign = cast(Assign, code.body[0])
self.result = cast(Name, assign.targets[0]).id
params = set()
variables = set()
constants = set()
op = None
for node in walk(assign.value):
if isinstance(node, Name):
name = node.id
if name.isupper():
variables.add(name)
else:
params.add(name)
elif isinstance(node, Constant):
constants.add(node.value)
elif isinstance(node, BinOp):
op = node.op
self.left = self.__to_name(node.left)
self.right = self.__to_name(node.right)
if op is None and len(constants) == 1:
self.left = next(iter(constants))
self.right = None
self.operator = op
self.parameters = frozenset(params)
self.variables = frozenset(variables)
self.constants = frozenset(constants)
self.compiled = compile(self.code, "", mode="exec")
def __to_name(self, node):
if isinstance(node, Name):
return node.id
elif isinstance(node, Constant):
return node.value
else:
return None
def __to_opsymbol(self, op):
if isinstance(op, Mult):
return "*"
elif isinstance(op, Div):
return "/"
elif isinstance(op, Add):
return "+"
elif isinstance(op, Sub):
return "-"
elif isinstance(op, Pow):
return "^"
return ""
def __str__(self):
return f"{self.result} = {self.left}{self.__to_opsymbol(self.operator)}{self.right}"
def __repr__(self):
return f"CodeOp({self.result} = f(params={self.parameters}, vars={self.variables}, consts={self.constants}))"
def __call__(self, *args, **kwargs: Mod) -> Mod:
"""Execute this operation with kwargs."""
loc = dict(kwargs)
exec(self.compiled, {}, loc)
return loc[self.result]
|
