def visit_Compare(self, node): node = self.generic_visit(node) if not all(self._has_matching_func(op) for op in node.ops): if len(node.ops) == 1: # Basic expressions are safe to leave as they are. return node else: raise NotImplementedError( 'compound expression with at least one unsupported ' 'operator: {}'.format(node.ops)) ops_and_comps = list(zip(node.ops, node.comparators)) left = node.left op_tree = None # Repeated comparisons are converted to conjunctions: # a < b < c -> a < b and b < c while ops_and_comps: op, right = ops_and_comps.pop(0) binary_comparison = self._as_function( self._matching_func(op), (left, right)) if isinstance(left, gast.Name) and isinstance(right, gast.Name): anno.setanno(binary_comparison, SAFE_BOOLEAN_OPERAND, True) if op_tree: self._expect_simple_symbol(right) op_tree = self._as_function('tf.logical_and', (binary_comparison, op_tree)) else: op_tree = binary_comparison left = right assert op_tree is not None return op_tree
def _track_symbol(self, node, composite_writes_alter_parent=False, writes_create_symbol=False): # A QN may be missing when we have an attribute (or subscript) on a function # call. Example: a().b if not anno.hasanno(node, anno.Basic.QN): return qn = anno.getanno(node, anno.Basic.QN) if isinstance(node.ctx, gast.Store): self.scope.mark_write(qn) if qn.is_composite and composite_writes_alter_parent: self.scope.mark_write(qn.parent) if writes_create_symbol: self.scope.mark_creation(qn, writes_create_symbol=True) if self._in_aug_assign: self.scope.mark_read(qn) elif isinstance(node.ctx, gast.Load): self.scope.mark_read(qn) elif isinstance(node.ctx, gast.Param): # Param contexts appear in function defs, so they have the meaning of # defining a variable. self.scope.mark_write(qn) self.scope.mark_param(qn, self.enclosing_entities[-1]) else: raise ValueError('Unknown context %s for node %s.' % (type(node.ctx), qn)) anno.setanno(node, NodeAnno.IS_LOCAL, self.scope.has(qn)) if self._in_return_statement: self.scope.mark_returned(qn)
def visit_Compare(self, node): node = self.generic_visit(node) ops_and_comps = list(zip(node.ops, node.comparators)) left = node.left op_tree = None # Repeated comparisons are converted to conjunctions: # a < b < c -> a < b and b < c while ops_and_comps: op, right = ops_and_comps.pop(0) binary_comparison = self._as_function(self._matching_func(op), (left, right)) if isinstance(left, gast.Name) and isinstance(right, gast.Name): anno.setanno(binary_comparison, SAFE_BOOLEAN_OPERAND, True) if op_tree: self._expect_simple_symbol(right) op_tree = self._as_function('ag__.and_', (op_tree, binary_comparison), args_as_lambda=True) else: op_tree = binary_comparison left = right assert op_tree is not None return op_tree
def visit_Lambda(self, node): assert not self._in_function_def_args self.state[_Lambda].enter() node = self.generic_visit(node) anno.setanno(node, anno.Static.SCOPE, self.scope) self.state[_Lambda].exit() return node
def visit_Attribute(self, node): parent_types = self.visit(node.value) # Attempt to use the static value if known. parent_value = anno.Static.VALUE.of(node.value, None) if parent_value is not None: static_value = getattr(parent_value, node.attr, None) else: # Fall back to the type if that is known. if parent_types is None: return None inferred_values = [ getattr(t, node.attr, None) for t in parent_types ] if not inferred_values: return None static_value = inferred_values[0] if static_value is None: return None if any(v is not static_value for v in inferred_values[1:]): # Static value not stable, assume it's dynamic. return None types = self.resolver.res_value(self.namespace, static_value) anno.setanno(node, anno.Static.VALUE, static_value) if __debug__: self._check_set(types) return types
def visit_Attribute(self, node): node = self.generic_visit(node) parent_val = anno.getanno(node.value, "static_value", default=None) if parent_val is not None: if hasattr(parent_val, node.attr): anno.setanno(node, "static_value", getattr(parent_val, node.attr)) return node
def visit_With(self, node): node.items = self.visit_block(node.items) node.body, definitely_returns = self._visit_statement_block( node, node.body) if definitely_returns: anno.setanno(node, STMT_DEFINITELY_RETURNS, True) return node
def _block_statement_live_out(self, node): successors = self.current_analyzer.graph.stmt_next[node] stmt_live_out = set() for s in successors: stmt_live_out.update(self.current_analyzer.in_[s]) anno.setanno(node, anno.Static.LIVE_VARS_OUT, frozenset(stmt_live_out)) return node
def test_create_source_map_multiple_nodes(self): source = """ from __future__ import print_function def test_fn(x): return x + 1 """ source = textwrap.dedent(source) nodes = parser.parse_str(source, single_node=False) fake_import_origin = origin_info.OriginInfo( loc=origin_info.Location('fake_filename', 3, 7), function_name='fake_function_name', source_code_line='fake source line', comment=None) anno.setanno(nodes[0], anno.Basic.ORIGIN, fake_import_origin) fake_function_origin = origin_info.OriginInfo( loc=origin_info.Location('fake_filename', 3, 7), function_name='fake_function_name', source_code_line='fake source line', comment=None) anno.setanno(nodes[1], anno.Basic.ORIGIN, fake_function_origin) source_map = origin_info.create_source_map(nodes, source, 'test_filename') loc = origin_info.LineLocation('test_filename', 2) self.assertIn(loc, source_map) self.assertIs(source_map[loc], fake_import_origin) loc = origin_info.LineLocation('test_filename', 3) self.assertIn(loc, source_map) self.assertIs(source_map[loc], fake_function_origin)
def visit_Lambda(self, node): with self.state[_Function] as fn_scope: node = self.generic_visit(node) # Only wrap the top-level function. Theoretically, we can and should wrap # everything, but that can lead to excessive boilerplate when lambdas are # nested. # TODO(mdan): Looks more closely for use cases that actually require this. if fn_scope.level > 2: return templates.replace_as_expression( 'ag__.autograph_artifact(l)', l=node) scope = anno.getanno(node, anno.Static.SCOPE) function_context_name = self.ctx.namer.new_symbol( 'lscope', scope.referenced) fn_scope.context_name = function_context_name anno.setanno(node, 'function_context_name', function_context_name) template = """ ag__.with_function_scope( lambda function_context: body, function_context_name, options) """ node.body = templates.replace_as_expression( template, options=self._function_scope_options(fn_scope).to_ast(), function_context=function_context_name, function_context_name=gast.Constant(function_context_name, kind=None), body=node.body) return node
def _block_statement_live_out(self, node): successors = self.current_analyzer.graph.stmt_next[node] stmt_live_out = set() for s in successors: stmt_live_out.update(self.current_analyzer.in_[s]) anno.setanno(node, anno.Static.LIVE_VARS_OUT, frozenset(stmt_live_out)) return node
def visit_Attribute(self, node): node = self.generic_visit(node) parent_val = anno.getanno(node.value, STATIC_VALUE, default=None) if parent_val is not None and tf_inspect.ismodule(parent_val): if hasattr(parent_val, node.attr): anno.setanno(node, STATIC_VALUE, getattr(parent_val, node.attr)) return node
def _track_symbol(self, node, composite_writes_alter_parent=False, writes_create_symbol=False): # A QN may be missing when we have an attribute (or subscript) on a function # call. Example: a().b if not anno.hasanno(node, anno.Basic.QN): return qn = anno.getanno(node, anno.Basic.QN) if isinstance(node.ctx, gast.Store): self.scope.mark_write(qn) if qn.is_composite and composite_writes_alter_parent: self.scope.mark_write(qn.parent) if writes_create_symbol: self.scope.mark_creation(qn, writes_create_symbol=True) if self._in_aug_assign: self.scope.mark_read(qn) elif isinstance(node.ctx, gast.Load): self.scope.mark_read(qn) elif isinstance(node.ctx, gast.Param): # Param contexts appear in function defs, so they have the meaning of # defining a variable. self.scope.mark_write(qn) self.scope.mark_param(qn, self.enclosing_entities[-1]) else: raise ValueError('Unknown context %s for node %s.' % (type(node.ctx), qn)) anno.setanno(node, NodeAnno.IS_LOCAL, self.scope.has(qn)) if self._in_return_statement: self.scope.mark_returned(qn)
def visit_FunctionDef(self, node): # The FunctionDef node itself has a Scope object that tracks the creation # of its name, along with the usage of any decorator accompany it. self._enter_scope(False) node.decorator_list = self.visit_block(node.decorator_list) self.scope.mark_modified(qual_names.QN(node.name)) anno.setanno(node, anno.Static.SCOPE, self.scope) self._exit_scope() # A separate Scope tracks the actual function definition. self._enter_scope(True) assert not (self._in_function_def_args or self.state[_Lambda].level) self._in_function_def_args = True node.args = self.visit(node.args) self._in_function_def_args = False # Track the body separately. This is for compatibility reasons, it may not # be strictly needed. self._enter_scope(False) node.body = self.visit_block(node.body) anno.setanno(node, NodeAnno.BODY_SCOPE, self.scope) self._exit_scope() self._exit_scope() return node
def visit_Attribute(self, node): node = self.generic_visit(node) parent_val = anno.getanno(node.value, STATIC_VALUE, default=None) if parent_val is not None and tf_inspect.ismodule(parent_val): if hasattr(parent_val, node.attr): anno.setanno(node, STATIC_VALUE, getattr(parent_val, node.attr)) return node
def test_create_source_map_multiple_nodes(self): source = """ from __future__ import print_function def test_fn(x): return x + 1 """ source = textwrap.dedent(source) nodes = parser.parse_str(source, single_node=False) fake_import_origin = origin_info.OriginInfo( loc=origin_info.Location('fake_filename', 3, 7), function_name='fake_function_name', source_code_line='fake source line', comment=None) anno.setanno(nodes[0], anno.Basic.ORIGIN, fake_import_origin) fake_function_origin = origin_info.OriginInfo( loc=origin_info.Location('fake_filename', 3, 7), function_name='fake_function_name', source_code_line='fake source line', comment=None) anno.setanno(nodes[1], anno.Basic.ORIGIN, fake_function_origin) source_map = origin_info.create_source_map(nodes, source, 'test_filename') loc = origin_info.LineLocation('test_filename', 2) self.assertIn(loc, source_map) self.assertIs(source_map[loc], fake_import_origin) loc = origin_info.LineLocation('test_filename', 3) self.assertIn(loc, source_map) self.assertIs(source_map[loc], fake_function_origin)
def visit_Compare(self, node): node = self.generic_visit(node) if not all(self._has_matching_func(op) for op in node.ops): if len(node.ops) == 1: # Basic expressions are safe to leave as they are. return node else: raise NotImplementedError( 'compound expression with at least one unsupported ' 'operator: {}'.format(node.ops)) ops_and_comps = list(zip(node.ops, node.comparators)) left = node.left op_tree = None # Repeated comparisons are converted to conjunctions: # a < b < c -> a < b and b < c while ops_and_comps: op, right = ops_and_comps.pop(0) binary_comparison = self._as_function(self._matching_func(op), (left, right)) if isinstance(left, gast.Name) and isinstance(right, gast.Name): anno.setanno(binary_comparison, SAFE_BOOLEAN_OPERAND, True) if op_tree: self._expect_simple_symbol(right) op_tree = self._as_function('tf.logical_and', (binary_comparison, op_tree)) else: op_tree = binary_comparison left = right assert op_tree is not None return op_tree
def visit_Print(self, node): self._enter_scope(False) node.values = self.visit_block(node.values) anno.setanno(node, anno.Static.SCOPE, self.scope) anno.setanno(node, NodeAnno.ARGS_SCOPE, self.scope) self._exit_scope() return node
def visit_FunctionDef(self, node): # The FunctionDef node itself has a Scope object that tracks the creation # of its name, along with the usage of any decorator accompany it. self._enter_scope(False) node.decorator_list = self.visit_block(node.decorator_list) self.scope.mark_modified(qual_names.QN(node.name)) anno.setanno(node, anno.Static.SCOPE, self.scope) self._exit_scope() # A separate Scope tracks the actual function definition. self._enter_scope(True) assert not (self._in_function_def_args or self.state[_Lambda].level) self._in_function_def_args = True node.args = self.visit(node.args) self._in_function_def_args = False # Track the body separately. This is for compatibility reasons, it may not # be strictly needed. self._enter_scope(False) node.body = self.visit_block(node.body) anno.setanno(node, NodeAnno.BODY_SCOPE, self.scope) self._exit_scope() self._exit_scope() return node
def visit_Lambda(self, node): with self.state[_Function] as fn_scope: node = self.generic_visit(node) # TODO(mdan): Fix the tests so that we can always add this decorator. if fn_scope.level > 2: return templates.replace_as_expression( 'ag__.autograph_artifact(l)', l=node) scope = anno.getanno(node, anno.Static.SCOPE) function_context_name = self.ctx.namer.new_symbol('lscope', scope.referenced) fn_scope.context_name = function_context_name anno.setanno(node, 'function_context_name', function_context_name) template = """ ag__.with_function_scope( lambda function_context: body, function_context_name, options) """ node.body = templates.replace_as_expression( template, options=self._function_scope_options(fn_scope).to_ast(), function_context=function_context_name, function_context_name=gast.Constant(function_context_name, kind=None), body=node.body) return node
def visit_Print(self, node): self._enter_scope(False) node.values = self.visit_block(node.values) anno.setanno(node, anno.Static.SCOPE, self.scope) anno.setanno(node, NodeAnno.ARGS_SCOPE, self.scope) self._exit_scope() return node
def visit_Attribute(self, node): node = self.generic_visit(node) if anno.hasanno(node.value, anno.Basic.QN): anno.setanno( node, anno.Basic.QN, QN(anno.getanno(node.value, anno.Basic.QN), attr=node.attr)) return node
def visit_For(self, node): original_node = node scope = anno.getanno(node, NodeAnno.BODY_SCOPE) break_var = self.ctx.namer.new_symbol('break_', scope.referenced) node.target = self.visit(node.target) node.iter = self.visit(node.iter) node.body, break_used = self._process_body(node.body, break_var) # A break in the else clause applies to the containing scope. node.orelse = self.visit_block(node.orelse) if not break_used: template = """ for target in iter_: body orelse """ node = templates.replace( template, iter_=node.iter, target=node.target, body=node.body, orelse=node.orelse) new_for_node = node[0] anno.copyanno(original_node, new_for_node, anno.Basic.EXTRA_LOOP_TEST) anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES) return node # Python's else clause only triggers if the loop exited cleanly (e.g. # break did not trigger). guarded_orelse = self._guard_if_present(node.orelse, break_var) extra_test = templates.replace_as_expression( 'ag__.not_(var_name)', var_name=break_var) # The extra test is hidden in the AST, which will confuse the static # analysis. To mitigate that, we insert a no-op statement that ensures # the control variable is marked as used. # TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name) template = """ var_name = False for target in iter_: (var_name,) body orelse """ node = templates.replace( template, var_name=break_var, iter_=node.iter, target=node.target, body=node.body, orelse=guarded_orelse) new_for_node = node[1] anno.setanno(new_for_node, anno.Basic.EXTRA_LOOP_TEST, extra_test) anno.copyanno(original_node, new_for_node, anno.Basic.DIRECTIVES) return node
def _aggregate_predecessors_defined_in(self, node): preds = self.current_analyzer.graph.stmt_prev[node] node_defined_in = set() for p in preds: node_defined_in |= set(self.current_analyzer.out[p].value.keys()) anno.setanno(node, anno.Static.DEFINED_VARS_IN, frozenset(node_defined_in))
def visit_Lambda(self, node): self.state[_Function].enter() node = self.generic_visit(node) # Only wrap the top-level function. Theoretically, we can and should wrap # everything, but that can lead to excessive boilerplate when lambdas are # nested. # TODO(mdan): Looks more closely for use cases that actually require this. if self.state[_Function].level > 2: self.state[_Function].exit() return node scope = anno.getanno(node, anno.Static.SCOPE) function_context_name = self.ctx.namer.new_symbol( 'lscope', scope.referenced) self.state[_Function].context_name = function_context_name anno.setanno(node, 'function_context_name', function_context_name) template = """ ag__.with_function_scope( lambda function_context: body, function_context_name, options) """ node.body = templates.replace_as_expression( template, options=self.ctx.program.options.to_ast(), function_context=function_context_name, function_context_name=gast.Str(function_context_name), body=node.body) self.state[_Function].exit() return node
def visit_Name(self, node): node = self.generic_visit(node) if isinstance(node.ctx, gast.Load) and node.id in self.ctx.info.namespace: anno.setanno(node, "static_value", self.ctx.info.namespace[node.id]) return node
def _as_function(self, func_name, args): template = """ func_name(args) """ replacement = templates.replace_as_expression( template, func_name=parser.parse_expression(func_name), args=args) anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True) return replacement
def visit(self, node): types = super().visit(node) if __debug__: self._check_set(types) if types is not None: # TODO(mdan): Normalize by removing subtypes. anno.setanno(node, anno.Static.TYPES, tuple(types)) return types
def _aggregate_successors_live_in(self, node): successors = self.current_analyzer.graph.stmt_next[node] node_live_out = set() for s in successors: node_live_out.update(self.current_analyzer.in_[s]) anno.setanno(node, anno.Static.LIVE_VARS_OUT, frozenset(node_live_out)) node = self.generic_visit(node) return node
def visit(self, node): node = super(Annotator, self).visit(node) if (self.current_analyzer is not None and isinstance(node, gast.stmt) and node in self.current_analyzer.graph.index): cfg_node = self.current_analyzer.graph.index[node] anno.setanno(node, anno.Static.LIVE_VARS_IN, frozenset(self.current_analyzer.in_[cfg_node])) return node
def visit_Name(self, node): node = self.generic_visit(node) if isinstance(node.ctx, gast.Load): defs = anno.getanno(node, anno.Static.DEFINITIONS, ()) is_defined = bool(defs) if not is_defined and node.id in self.ctx.info.namespace: anno.setanno(node, STATIC_VALUE, self.ctx.info.namespace[node.id]) return node
def _aggregate_successors_live_in(self, node): successors = self.current_analyzer.graph.stmt_next[node] node_live_out = set() for s in successors: node_live_out.update(self.current_analyzer.in_[s]) anno.setanno(node, anno.Static.LIVE_VARS_OUT, frozenset(node_live_out)) node = self.generic_visit(node) return node
def _as_function(self, func_name, args): template = """ func_name(args) """ replacement = templates.replace_as_expression( template, func_name=parser.parse_expression(func_name), args=args) anno.setanno(replacement, SAFE_BOOLEAN_OPERAND, True) return replacement
def visit(self, node): if not isinstance(node, gast.AST): # This is not that uncommon a mistake: various node bodies are lists, for # example, posing a land mine for transformers that need to recursively # call `visit`. The error needs to be raised before the exception handler # below is installed, because said handler will mess up if `node` is not, # in fact, a node. msg = ('invalid value for "node": expected "ast.AST", got "{}"; to' ' visit lists of nodes, use "visit_block" instead').format( type(node)) raise ValueError(msg) if anno.hasanno(node, anno.Basic.SKIP_PROCESSING): return node parent_origin = self.ctx.current_origin if anno.hasanno(node, anno.Basic.ORIGIN): self.ctx.current_origin = anno.getanno(node, anno.Basic.ORIGIN) try: processing_expr_node = isinstance(node, gast.Expr) if processing_expr_node: entry_expr_value = node.value result = super(Base, self).visit(node) # Adjust for consistency: replacing the value of an Expr with # an Assign node removes the need for the Expr node. if (processing_expr_node and isinstance(result, gast.Expr) and (result.value is not entry_expr_value)): # When the replacement is a list, it is assumed that the list came # from a template that contained a number of statements, which # themselves are standalone and don't require an enclosing Expr. if isinstance(result.value, (list, tuple, gast.Assign, gast.AugAssign)): result = result.value # By default, all replacements receive the origin info of the replaced # node. if result is not node and result is not None: inherited_origin = anno.getanno(node, anno.Basic.ORIGIN, default=parent_origin) if inherited_origin is not None: nodes_to_adjust = result if isinstance(result, (list, tuple)): nodes_to_adjust = result else: nodes_to_adjust = (result, ) for n in nodes_to_adjust: if not anno.hasanno(n, anno.Basic.ORIGIN): anno.setanno(n, anno.Basic.ORIGIN, inherited_origin) finally: self.ctx.current_origin = parent_origin return result
def visit_Call(self, node): self._enter_scope(False) node.args = self.visit_block(node.args) node.keywords = self.visit_block(node.keywords) # TODO(mdan): Account starargs, kwargs anno.setanno(node, NodeAnno.ARGS_SCOPE, self.scope) self._exit_scope() node.func = self.visit(node.func) return node
def visit(self, node): node = super(Annotator, self).visit(node) if (self.current_analyzer is not None and isinstance(node, gast.stmt) and node in self.current_analyzer.graph.index): cfg_node = self.current_analyzer.graph.index[node] anno.setanno(node, anno.Static.LIVE_VARS_IN, frozenset(self.current_analyzer.in_[cfg_node])) return node
def _process_statement_directive(self, call_node, directive): if self.local_scope_level < 1: raise ValueError( '"%s" must be used inside a statement' % directive.__name__) target = self.get_local(ENCLOSING_LOOP) node_anno = anno.getanno(target, converter.AgAnno.DIRECTIVES, {}) node_anno[directive] = _map_args(call_node, directive) anno.setanno(target, converter.AgAnno.DIRECTIVES, node_anno) return call_node
def visit_Call(self, node): self._enter_scope(False) node.args = self.visit_block(node.args) node.keywords = self.visit_block(node.keywords) # TODO(mdan): Account starargs, kwargs anno.setanno(node, NodeAnno.ARGS_SCOPE, self.scope) self._exit_scope() node.func = self.visit(node.func) return node
def test_rename_symbols_annotations(self): node = parser.parse_str('a[i]') node = qual_names.resolve(node) anno.setanno(node, 'foo', 'bar') orig_anno = anno.getanno(node, 'foo') node = ast_util.rename_symbols(node, {qual_names.QN('a'): qual_names.QN('b')}) self.assertIs(anno.getanno(node, 'foo'), orig_anno)
def visit_For(self, node): self._enter_scope(False) node.target = self.visit(node.target) node.iter = self.visit(node.iter) anno.setanno(node.iter, anno.Static.SCOPE, self.scope) self._exit_scope() node = self._process_parallel_blocks(node, ((node.body, NodeAnno.BODY_SCOPE), (node.orelse, NodeAnno.ORELSE_SCOPE))) return node
def visit_While(self, node): self._enter_scope(False) node.test = self.visit(node.test) anno.setanno(node, NodeAnno.COND_SCOPE, self.scope) anno.setanno(node.test, anno.Static.SCOPE, self.scope) self._exit_scope() node = self._process_parallel_blocks( node, ((node.body, NodeAnno.BODY_SCOPE), (node.orelse, NodeAnno.ORELSE_SCOPE))) return node
def test_copy(self): node_1 = ast.Name() anno.setanno(node_1, 'foo', 3) node_2 = ast.Name() anno.copyanno(node_1, node_2, 'foo') anno.copyanno(node_1, node_2, 'bar') self.assertTrue(anno.hasanno(node_2, 'foo')) self.assertFalse(anno.hasanno(node_2, 'bar'))
def visit_If(self, node): self._enter_scope(False) node.test = self.visit(node.test) node_scope = self._exit_and_record_scope(node.test) anno.setanno(node, NodeAnno.COND_SCOPE, node_scope) node = self._process_parallel_blocks( node, ((node.body, NodeAnno.BODY_SCOPE), (node.orelse, NodeAnno.ORELSE_SCOPE))) return node
def test_rename_symbols_annotations(self): node = parser.parse_str('a[i]') node = qual_names.resolve(node) anno.setanno(node, 'foo', 'bar') orig_anno = anno.getanno(node, 'foo') node = ast_util.rename_symbols( node, {qual_names.QN('a'): qual_names.QN('b')}) self.assertIs(anno.getanno(node, 'foo'), orig_anno)
def visit_While(self, node): self._enter_scope(False) node.test = self.visit(node.test) anno.setanno(node, NodeAnno.COND_SCOPE, self.scope) anno.setanno(node.test, anno.Static.SCOPE, self.scope) self._exit_scope() node = self._process_parallel_blocks(node, ((node.body, NodeAnno.BODY_SCOPE), (node.orelse, NodeAnno.ORELSE_SCOPE))) return node
def resolve(nodes, source, function=None): """Adds an origin information to all nodes inside the body of function. Args: nodes: Union[ast.AST, Iterable[ast.AST, ...]] source: Text, the source code string for the function whose body nodes will be annotated. function: Callable, the function that will have all nodes inside of it annotation with an OriginInfo annotation with key anno.Basic.ORIGIN. If it is None then only the line numbers and column offset will be set in the annotation, with the rest of the information being None. Returns: A tuple of the AST node for function and a String containing its source code. """ if not isinstance(nodes, (list, tuple)): nodes = (nodes,) if function: _, function_lineno = tf_inspect.getsourcelines(function) function_filepath = tf_inspect.getsourcefile(function) else: function_lineno = None function_filepath = None # TODO(mdan): Pull this to a separate utility. code_reader = six.StringIO(source) comment_map = {} for token in tokenize.generate_tokens(code_reader.readline): tok_type, tok_string, loc, _, _ = token srow, _ = loc if tok_type == tokenize.COMMENT: comment_map[srow] = tok_string.strip()[1:].strip() source_lines = source.split('\n') for node in nodes: for n in gast.walk(node): if not hasattr(n, 'lineno'): continue lineno_in_body = n.lineno source_code_line = source_lines[lineno_in_body - 1] if function: source_lineno = function_lineno + lineno_in_body function_name = function.__name__ else: source_lineno = lineno_in_body function_name = None location = Location(function_filepath, source_lineno, n.col_offset) origin = OriginInfo(location, function_name, source_code_line, comment_map.get(source_lineno)) anno.setanno(n, anno.Basic.ORIGIN, origin)
def _block_statement_live_in(self, node, entry_node): if entry_node in self.current_analyzer.graph.index: cfg_node = self.current_analyzer.graph.index[entry_node] stmt_live_in = frozenset(self.current_analyzer.in_[cfg_node]) else: assert anno.hasanno(entry_node, anno.Static.LIVE_VARS_IN), ( 'If not matching a CFG node, must be a block statement:' ' {}'.format(entry_node)) stmt_live_in = anno.getanno(entry_node, anno.Static.LIVE_VARS_IN) anno.setanno(node, anno.Static.LIVE_VARS_IN, stmt_live_in) return node
def test_copy_clean_preserves_annotations(self): node = parser.parse_str( textwrap.dedent(""" def f(a): return a + 1 """)) anno.setanno(node.body[0], 'foo', 'bar') anno.setanno(node.body[0], 'baz', 1) new_node = ast_util.copy_clean(node, preserve_annos={'foo'}) self.assertEqual(anno.getanno(new_node.body[0], 'foo'), 'bar') self.assertFalse(anno.hasanno(new_node.body[0], 'baz'))
def _process_function_arg(self, arg_node): qn = anno.getanno(arg_node, anno.Basic.QN) arg_name = str(qn) self.scope.setval(qn, arg_node) if (len(self.enclosing_entities) == 1 and arg_name in self.entity_info.arg_types): # Forge a node to hold the type information, so that method calls on # it can resolve the type. type_string, type_obj = self.entity_info.arg_types[arg_name] anno.setanno(arg_node, 'type', type_obj) anno.setanno(arg_node, 'type_fqn', tuple(type_string.split('.')))
def resolve(node, source, function=None): """Adds an origin information to node and its subnodes. This allows us to map the original source code line numbers to generated source code. Args: node: gast.AST node. Should be a gast.FunctionDef. This is the node we annotate with origin information. source: Text, the source code. Should satisfy relationship `node in iter_tree(gast.parse(source))`; otherwise the lineno will be unreliable. function: The original function. If it is None then only the line numbers and column offset will be set in the annotation, with the rest of the information being None. """ if function: _, function_lineno = tf_inspect.getsourcelines(function) function_filepath = tf_inspect.getsourcefile(function) else: function_lineno = None function_filepath = None # TODO(mdan): Pull this to a separate utility. code_reader = six.StringIO(source) comment_map = {} for token in tokenize.generate_tokens(code_reader.readline): tok_type, tok_string, loc, _, _ = token srow, _ = loc if tok_type == tokenize.COMMENT: comment_map[srow] = tok_string.strip()[1:].strip() source_lines = source.split('\n') for n in gast.walk(node): if not hasattr(n, 'lineno'): continue within_body_offset = n.lineno - node.lineno source_code_line = source_lines[n.lineno - 1] if function: source_lineno = function_lineno + within_body_offset function_name = function.__name__ else: source_lineno = n.lineno function_name = None location = Location(function_filepath, source_lineno, n.col_offset) origin = OriginInfo(location, function_name, source_code_line, comment_map.get(source_lineno)) anno.setanno(n, anno.Basic.ORIGIN, origin)
def visit_Call(self, node): # If the function call is wrapped by one of the marker decorators, # consider it graph ready. if anno.hasanno(node.func, 'live_val'): target_entity = anno.getanno(node.func, 'live_val') if target_entity in self.ctx.program.options.strip_decorators: if len(node.args) < 1: raise ValueError( 'Found call to decorator function "%s", but it had no arguments. ' 'A decorator needs at least one positional argument.' % target_entity) anno.setanno(node.args[0], 'graph_ready', True) self.generic_visit(node) if anno.hasanno(node.func, 'live_val'): target_entity = anno.getanno(node.func, 'live_val') if anno.hasanno(node.func, 'fqn'): target_fqn = anno.getanno(node.func, 'fqn') else: target_fqn = None if self._function_is_compilable(target_entity): node = self._rename_compilable_function(node) elif target_fqn and target_fqn in KNOWN_NUMPY_FUNCTIONS: # TODO(mdan): Should we replace these with equivalent TF ops instead? node = self._wrap_to_py_func_single_return( node, KNOWN_NUMPY_FUNCTIONS[target_fqn].dtype) else: raise NotImplementedError( 'py_func with return values (unknown function)') else: if anno.hasanno(node.func, anno.Basic.QN): # Special-case a few builtins that otherwise go undetected. This # normally doesn't pose a problem, but the dict built-in doesn't # work with inspect.getargspec which is required for dynamic functions. # Note: expecting this is resilient to aliasing (e.g. # dict = an_evil_dict), because in those cases the regular mechanisms # process a simple user function. qn = anno.getanno(node.func, anno.Basic.QN) # Add items to this list as needed. if str(qn) in ('dict',): return node if ast_util.matches(node, 'super(_)'): # super() calls are preserved. The class conversion mechanism will # ensure that they return the correct value. return node if self.ctx.program.options.recursive: node = self._insert_dynamic_conversion(node) return node
def test_basic(self): def test_fn(): raise ValueError() node, ctx = self.prepare(test_fn, {}) anno.setanno( node, anno.Basic.ORIGIN, origin_info.OriginInfo(None, 'test_function_name', 'test_code', 'test_comment')) node = error_handlers.transform(node, ctx) with self.compiled(node, {}) as result: with self.assertRaises(errors.GraphConstructionError): # Here we just assert that the handler works. Its correctness is # verified by errors_test.py. result.test_fn()
def visit_If(self, node): node.test = self.visit(node.test) node.body, body_definitely_returns = self._visit_statement_block( node, node.body) if body_definitely_returns: anno.setanno(node, BODY_DEFINITELY_RETURNS, True) node.orelse, orelse_definitely_returns = self._visit_statement_block( node, node.orelse) if orelse_definitely_returns: anno.setanno(node, ORELSE_DEFINITELY_RETURNS, True) if body_definitely_returns and orelse_definitely_returns: self.state[_Block].definitely_returns = True return node
def visit_Name(self, node): if self.current_analyzer is None: # Names may appear outside function defs - for example in class # definitions. return node analyzer = self.current_analyzer cfg_node = self.current_cfg_node assert cfg_node is not None, 'name node outside of any statement?' qn = anno.getanno(node, anno.Basic.QN) if isinstance(node.ctx, gast.Load): anno.setanno(node, anno.Static.DEFINITIONS, tuple(analyzer.in_[cfg_node].value.get(qn, ()))) else: anno.setanno(node, anno.Static.DEFINITIONS, tuple(analyzer.out[cfg_node].value.get(qn, ()))) return node
def test_create_source_map(self): source = """ def test_fn(x): return x + 1 """ source = textwrap.dedent(source) node = parser.parse_str(source) fake_origin = origin_info.OriginInfo( loc=origin_info.Location('fake_filename', 3, 7), function_name='fake_function_name', source_code_line='fake source line', comment=None) anno.setanno(node, anno.Basic.ORIGIN, fake_origin) source_map = origin_info.create_source_map(node, source, 'test_filename') loc = origin_info.LineLocation('test_filename', 2) self.assertIn(loc, source_map) self.assertIs(source_map[loc], fake_origin)
def test_create_source_map(self): def test_fn(x): return x + 1 node, _ = parser.parse_entity(test_fn) fake_origin = origin_info.OriginInfo( loc=origin_info.Location('fake_filename', 3, 7), function_name='fake_function_name', source_code_line='fake source line', comment=None) fn_node = node.body[0] anno.setanno(fn_node.body[0], anno.Basic.ORIGIN, fake_origin) converted_code = compiler.ast_to_source(fn_node) source_map = origin_info.create_source_map( fn_node, converted_code, 'test_filename', [0]) loc = origin_info.LineLocation('test_filename', 2) self.assertIn(loc, source_map) self.assertIs(source_map[loc], fake_origin)
def visit_For(self, node): node.iter = self.visit(node.iter) node.target = self.visit(node.target) # Add the check for return to the loop condition. node.body = self._visit_statement_block(node, node.body) if self.state[_Return].used: extra_test = anno.getanno(node, 'extra_test', default=None) if extra_test is not None: extra_test = templates.replace_as_expression( 'ag__.and_(lambda: ag__.not_(control_var), lambda: extra_test)', extra_test=extra_test, control_var=self.state[_Function].do_return_var_name) else: extra_test = templates.replace_as_expression( 'ag__.not_(control_var)', control_var=self.state[_Function].do_return_var_name) anno.setanno(node, 'extra_test', extra_test) node.orelse = self._visit_statement_block(node, node.orelse) return node
def visit_For(self, node): scope = anno.getanno(node, NodeAnno.BODY_SCOPE) break_var = self.ctx.namer.new_symbol('break_', scope.referenced) node.target = self.visit(node.target) node.iter = self.visit(node.iter) node.body, break_used = self._process_body(node.body, break_var) # A break in the else clause applies to the containing scope. node.orelse = self.visit_block(node.orelse) if break_used: # Python's else clause only triggers if the loop exited cleanly (e.g. # break did not trigger). guarded_orelse = self._guard_if_present(node.orelse, break_var) extra_test = templates.replace_as_expression( 'not var_name', var_name=break_var) # The extra test is hidden in the AST, which will confuse the static # analysis. To mitigate that, we insert a no-op statement that ensures # the control variable is marked as used. # TODO(mdan): Use a marker instead, e.g. ag__.condition_loop_on(var_name) template = """ var_name = tf.constant(False) for target in iter_: (var_name,) body else: orelse """ node = templates.replace( template, var_name=break_var, iter_=node.iter, target=node.target, body=node.body, orelse=guarded_orelse) anno.setanno(node[1], 'extra_test', extra_test) return node