1# Copyright 2015 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"""Functions used to extract and analyze stacks. Faster than Python libs.""" 16# pylint: disable=g-bad-name 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21import linecache 22import sys 23 24# Names for indices into TF traceback tuples. 25TB_FILENAME = 0 26TB_LINENO = 1 27TB_FUNCNAME = 2 28TB_CODEDICT = 3 # Dictionary of Python interpreter state. 29 30 31def extract_stack(extract_frame_info_fn=None): 32 """A lightweight, extensible re-implementation of traceback.extract_stack. 33 34 NOTE(mrry): traceback.extract_stack eagerly retrieves the line of code for 35 each stack frame using linecache, which results in an abundance of stat() 36 calls. This implementation does not retrieve the code, and any consumer 37 should apply _convert_stack to the result to obtain a traceback that can 38 be formatted etc. using traceback methods. 39 40 Args: 41 extract_frame_info_fn: Optional callable fn(stack_frame) applied to each 42 stack frame. This callable's return value is stored as the sixth (last) 43 element of the returned tuples. If not provided, the returned tuples 44 will have None as their sixth value. 45 46 Returns: 47 A list of 6-tuples 48 (filename, lineno, name, frame_globals, func_start_lineno, custom_info) 49 corresponding to the call stack of the current thread. The returned tuples 50 have the innermost stack frame at the end, unlike the Python inspect 51 module's stack() function. 52 """ 53 default_fn = lambda f: None 54 extract_frame_info_fn = extract_frame_info_fn or default_fn 55 try: 56 raise ZeroDivisionError 57 except ZeroDivisionError: 58 f = sys.exc_info()[2].tb_frame.f_back 59 ret = [] 60 while f is not None: 61 lineno = f.f_lineno 62 co = f.f_code 63 filename = co.co_filename 64 name = co.co_name 65 frame_globals = f.f_globals 66 func_start_lineno = co.co_firstlineno 67 frame_info = extract_frame_info_fn(f) 68 ret.append((filename, lineno, name, frame_globals, func_start_lineno, 69 frame_info)) 70 f = f.f_back 71 ret.reverse() 72 return ret 73 74 75def convert_stack(stack, include_func_start_lineno=False): 76 """Converts a stack extracted using extract_stack() to a traceback stack. 77 78 Args: 79 stack: A list of n 5-tuples, 80 (filename, lineno, name, frame_globals, func_start_lineno). 81 include_func_start_lineno: True if function start line number should be 82 included as the 5th entry in return tuples. 83 84 Returns: 85 A list of n 4-tuples or 5-tuples 86 (filename, lineno, name, code, [optional: func_start_lineno]), where the 87 code tuple element is calculated from the corresponding elements of the 88 input tuple. 89 """ 90 ret = [] 91 for (filename, lineno, name, frame_globals, func_start_lineno, 92 unused_frame_info) in stack: 93 linecache.checkcache(filename) 94 line = linecache.getline(filename, lineno, frame_globals) 95 if line: 96 line = line.strip() 97 else: 98 line = None 99 if include_func_start_lineno: 100 ret.append((filename, lineno, name, line, func_start_lineno)) 101 else: 102 ret.append((filename, lineno, name, line)) 103 return ret 104