Skip to content

Commit

Permalink
Merge pull request #53 from sola-st/aryaz/annotations
Browse files Browse the repository at this point in the history
Fixed issue with annotated assignments
  • Loading branch information
AryazE authored Mar 10, 2024
2 parents 5b98c59 + de1609d commit 7aa6c43
Show file tree
Hide file tree
Showing 3 changed files with 28 additions and 10 deletions.
27 changes: 17 additions & 10 deletions src/dynapyt/instrument/CodeInstrumenter.py
Original file line number Diff line number Diff line change
Expand Up @@ -38,12 +38,16 @@ def __init__(self, src, file_path, iids: IIDs, selected_hooks):
self.current_function = []
self.selected_hooks = {
hook: {
"only": [re.compile(p) for p in details["only"]]
if "only" in details
else [],
"ignore": [re.compile(p) for p in details["ignore"]]
if "ignore" in details
else [],
"only": (
[re.compile(p) for p in details["only"]]
if "only" in details
else []
),
"ignore": (
[re.compile(p) for p in details["ignore"]]
if "ignore" in details
else []
),
}
for hook, details in selected_hooks.items()
}
Expand Down Expand Up @@ -1037,7 +1041,7 @@ def leave_AnnAssign(self, original_node, updated_node):
self.blacklist_nodes.append(cst.SimpleStatementLine(body=[original_node]))
self.blacklist_nodes.append(cst.Newline(value="\n"))
if "write" not in self.selected_hooks or original_node.value is None:
return updated_node
return original_node
callee_name = cst.Attribute(
value=cst.Name(value="_rt"), attr=cst.Name(value="_write_")
)
Expand All @@ -1053,18 +1057,21 @@ def leave_AnnAssign(self, original_node, updated_node):
value=cst.List(
elements=[
cst.Element(
self.__wrap_in_lambda(original_node.target, updated_node.target)
self.__wrap_in_lambda(
original_node.target, original_node.target
)
)
]
)
)
call = cst.Call(func=callee_name, args=[ast_arg, iid_arg, val_arg, left_arg])
if m.matches(updated_node.value, m.Yield()):
return updated_node.with_changes(
value=updated_node.value.with_changes(value=call)
target=original_node.target,
value=updated_node.value.with_changes(value=call),
)
else:
return updated_node.with_changes(value=call)
return updated_node.with_changes(target=original_node.target, value=call)

def leave_AugAssign(self, original_node, updated_node):
if ("write" not in self.selected_hooks) and (
Expand Down
4 changes: 4 additions & 0 deletions tests/trace_single_hook/annotated_assignment/expected.txt
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,8 @@ begin execution
Writing 10 to [('a',)]
read value 10
Writing 10 to [('b',)]
read value <class 'trace_single_hook.annotated_assignment.program.X'>
Writing 10 to [('x',)]
Writing <trace_single_hook.annotated_assignment.program.X object at <...>> to [('c',)]
Writing hello to [('c', 'y')]
end execution
7 changes: 7 additions & 0 deletions tests/trace_single_hook/annotated_assignment/program.py
Original file line number Diff line number Diff line change
@@ -1,2 +1,9 @@
class X:
def __init__(self):
self.x: int = 10


a = 10
b: int = a
c = X()
c.y: str = "hello"

0 comments on commit 7aa6c43

Please sign in to comment.