def _get_replacement_args(self, args, keywords: bool):
replacement_args = []
for arg in args:
is_starred = isinstance(arg, ast.Starred)
is_kwstarred = keywords and arg.arg is None
if keywords or is_starred:
maybe_kwarg = getattr(arg, 'value')
else:
maybe_kwarg = arg
with fast.location_of(maybe_kwarg):
with self.attrsub_context(None):
visited_maybe_kwarg = self.visit(maybe_kwarg)
with self.attrsub_context(None):
new_arg_value = cast(ast.expr, fast.Call(
func=self._emitter_ast(),
args=[TraceEvent.argument.to_ast(), self._get_copy_id_ast(maybe_kwarg)],
keywords=fast.kwargs(
ret=visited_maybe_kwarg,
is_starred=fast.NameConstant(is_starred),
is_kwstarred=fast.NameConstant(is_kwstarred)
),
))
if keywords or is_starred:
setattr(arg, 'value', new_arg_value)
else:
arg = new_arg_value
replacement_args.append(arg)
return replacement_args
def visit_Dict(self, node: ast.Dict):
traced_keys: List[Optional[ast.expr]] = []
traced_values: List[ast.expr] = []
for k, v in zip(node.keys, node.values):
is_dict_unpack = (k is None)
if is_dict_unpack:
traced_keys.append(None)
else:
with fast.location_of(k):
traced_keys.append(fast.Call(
func=self._emitter_ast(),
args=[TraceEvent.dict_key.to_ast(), self._get_copy_id_ast(k)],
keywords=fast.kwargs(
ret=self.visit(k),
value_node_id=self._get_copy_id_ast(v),
dict_node_id=self._get_copy_id_ast(node),
)
))
with fast.location_of(v):
if is_dict_unpack:
key_node_id_ast = fast.NameConstant(None)
else:
key_node_id_ast = self._get_copy_id_ast(k)
traced_values.append(fast.Call(
func=self._emitter_ast(),
args=[TraceEvent.dict_value.to_ast(), self._get_copy_id_ast(v)],
keywords=fast.kwargs(
ret=self.visit(v),
key_node_id=key_node_id_ast,
dict_node_id=self._get_copy_id_ast(node),
)
))
node.keys = traced_keys
node.values = traced_values
return self.visit_literal(node, should_inner_visit=False)
def visit_List_or_Tuple(self, node: Union[ast.List, ast.Tuple]):
traced_elts: List[ast.expr] = []
is_load = isinstance(getattr(node, 'ctx', ast.Load()), ast.Load)
saw_starred = False
for i, elt in enumerate(node.elts):
if isinstance(elt, ast.Starred):
# TODO: trace starred elts too
saw_starred = True
traced_elts.append(elt)
continue
elif not is_load:
traced_elts.append(self.visit(elt))
continue
with fast.location_of(elt):
traced_elts.append(fast.Call(
func=self._emitter_ast(),
args=[
TraceEvent.list_elt.to_ast() if isinstance(node, ast.List) else TraceEvent.tuple_elt.to_ast(),
self._get_copy_id_ast(elt),
],
keywords=fast.kwargs(
ret=self.visit(elt),
index=fast.NameConstant(None) if saw_starred else fast.Num(i),
container_node_id=self._get_copy_id_ast(node),
)
))
node.elts = traced_elts
return self.visit_literal(node, should_inner_visit=False)
def visit_Attribute_or_Subscript(
self,
node: Union[ast.Attribute, ast.Subscript],
attr_or_sub: ast.expr,
call_context: bool = False
):
orig_node_id = id(node)
with fast.location_of(node.value):
extra_args: List[ast.keyword] = []
if isinstance(node.value, ast.Name):
extra_args = fast.kwargs(obj_name=fast.Str(node.value.id))
with self.attrsub_context(node):
node.value = fast.Call(
func=self._emitter_ast(),
args=[
TraceEvent.subscript.to_ast() if isinstance(node, ast.Subscript) else TraceEvent.attribute.to_ast(),
self._get_copy_id_ast(node.value)
],
keywords=fast.kwargs(
ret=self.visit(node.value),
attr_or_subscript=attr_or_sub,
ctx=fast.Str(node.ctx.__class__.__name__),
call_context=fast.NameConstant(call_context),
top_level_node_id=self._get_copy_id_ast(self._top_level_node_for_symbol)
) + extra_args
)
# end fast.location_of(node.value)
return self._maybe_wrap_symbol_in_before_after_tracing(node, orig_node_id=orig_node_id)
def _maybe_wrap_symbol_in_before_after_tracing(
self, node, call_context=False, orig_node_id=None, begin_kwargs=None, end_kwargs=None
):
if self._inside_attrsub_load_chain:
return node
orig_node = node
orig_node_id = orig_node_id or id(orig_node)
begin_kwargs = begin_kwargs or {}
end_kwargs = end_kwargs or {}
ctx = getattr(orig_node, 'ctx', ast.Load())
is_load = isinstance(ctx, ast.Load)
with fast.location_of(node):
begin_kwargs['ret'] = self._get_copy_id_ast(orig_node_id)
if is_load:
end_ret = orig_node
elif isinstance(orig_node, (ast.Attribute, ast.Subscript)):
end_ret = orig_node.value
else:
raise TypeError('Unsupported node type for before / after symbol tracing: %s', type(orig_node))
end_kwargs['ret'] = end_ret
end_kwargs['ctx'] = fast.Str(ctx.__class__.__name__)
end_kwargs['call_context'] = fast.NameConstant(call_context)
node = fast.Call(
func=self._emitter_ast(),
args=[
TraceEvent.after_complex_symbol.to_ast(),
fast.Call(
# this will return the node id
func=self._emitter_ast(),
args=[TraceEvent.before_complex_symbol.to_ast(), self._get_copy_id_ast(orig_node_id)],
keywords=fast.kwargs(**begin_kwargs),
)
],
keywords=fast.kwargs(**end_kwargs),
)
if not is_load:
if isinstance(orig_node, ast.Attribute):
node = fast.Attribute(
value=node,
attr=orig_node.attr,
)
elif isinstance(orig_node, ast.Subscript):
node = fast.Subscript(
value=node,
slice=orig_node.slice,
)
else:
logger.error(
'Symbol tracing stores unsupported for node %s with type %s', orig_node, type(orig_node)
)
assert False
node.ctx = ast.Store()
# end location_of(node)
return node