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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
|
"""
This module provides a class for a code operation.
"""
from ast import (
Module,
walk,
Name,
BinOp,
UnaryOp,
Constant,
Mult,
Div,
Add,
Sub,
Pow,
Assign,
operator as ast_operator,
unaryop as ast_unaryop,
USub,
)
from enum import Enum
from types import CodeType
from typing import FrozenSet, cast, Any, Optional, Union
from public import public
from .mod import Mod
@public
class OpType(Enum):
"""A type of binary and unary operators."""
Add = (2, "+")
Sub = (2, "-")
Neg = (1, "-")
Mult = (2, "*")
Div = (2, "/")
Inv = (1, "/")
Sqr = (1, "^")
Pow = (2, "^")
Id = (1, "")
def __init__(self, num_inputs: int, op_str: str):
self.num_inputs = num_inputs
self.op_str = op_str
@public
class CodeOp(object):
"""An operation that can be executed."""
result: str
"""The result variable of the operation (e.g. the `r` in `r = 2*a`)."""
parameters: FrozenSet[str]
"""The parameters used in the operation (e.g. `a`, `b`)."""
variables: FrozenSet[str]
"""The variables used in the operation (e.g. `X1`, `Z2`)."""
code: Module
"""The code of the operation."""
operator: OpType
"""The operator type that executes in the operation."""
compiled: CodeType
"""The compiled code of the operation."""
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: Optional[Union[ast_operator, ast_unaryop]] = None
self.left = None
self.right = 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)
if isinstance(assign.value, BinOp):
op = assign.value.op
self.left = self.__to_name(assign.value.left)
self.right = self.__to_name(assign.value.right)
elif isinstance(assign.value, UnaryOp):
op = assign.value.op
self.right = self.__to_name(assign.value.operand)
elif isinstance(assign.value, Name):
self.left = assign.value.id
elif isinstance(assign.value, Constant):
self.left = assign.value.value
self.operator = self.__to_op(op, self.left, self.right)
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_op(
self, op: Optional[Union[ast_operator, ast_unaryop]], left: Any, right: Any
) -> OpType:
if isinstance(op, Mult):
return OpType.Mult
elif isinstance(op, Div):
if left == 1:
return OpType.Inv
return OpType.Div
elif isinstance(op, Add):
return OpType.Add
elif isinstance(op, Sub):
return OpType.Sub
elif isinstance(op, USub):
return OpType.Neg
elif isinstance(op, Pow):
if right == 2:
return OpType.Sqr
return OpType.Pow
return OpType.Id
def __str__(self):
return f"{self.result} = {self.left if self.left is not None else ''}{self.operator.op_str}{self.right if self.right is not None else ''}"
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."""
exec(self.compiled, {}, kwargs)
return kwargs[self.result]
|
