• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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