1 # Copyright 2016 The TensorFlow Authors. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 # ============================================================================== 15 """Handles decorators. 16 17 Note: this module only deals with functions whose decorators are still recorded 18 in the AST. This does not always happen. See the unit test for an example. 19 """ 20 21 from __future__ import absolute_import 22 from __future__ import division 23 from __future__ import print_function 24 25 import gast 26 27 from tensorflow.contrib.py2tf.pyct import anno 28 from tensorflow.contrib.py2tf.pyct import pretty_printer 29 30 31 class DecoratorsTransformer(gast.NodeTransformer): 32 """Converts or removes decorators.""" 33 34 def __init__(self, remove_decorators): 35 self.remove_decorators = remove_decorators 36 37 # pylint:disable=invalid-name 38 39 def visit_FunctionDef(self, node): 40 self.generic_visit(node) 41 kept_decorators = [] 42 for dec in node.decorator_list: 43 if isinstance(dec, gast.Call): 44 dec_func = dec.func 45 else: 46 dec_func = dec 47 if not anno.hasanno(dec_func, 'live_val'): 48 raise ValueError( 49 'Could not resolve decorator: %s' % pretty_printer.fmt(dec_func)) 50 dec_value = anno.getanno(dec_func, 'live_val') 51 if dec_value not in self.remove_decorators: 52 kept_decorators.append(dec) 53 node.decorator_list = kept_decorators 54 return node 55 56 # pylint:enable=invalid-name 57 58 59 def transform(node, remove_decorators): 60 transformer = DecoratorsTransformer(remove_decorators) 61 node = transformer.visit(node) 62 return node 63
