Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/tensorflow/python/util/tf_stack.py: 47%

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 

17import collections 

18import inspect 

19import threading 

20 

21from tensorflow.python.util import _tf_stack 

22 

23# Generally such lookups should be done using `threading.local()`. See 

24# https://blogs.gnome.org/jamesh/2008/06/11/tls-python/ for a detailed 

25# explanation of why. However the transform stacks are expected to be empty 

26# when a thread is joined, so reusing the key does not introduce a correctness 

27# issue. Moreover, get_ident is faster than storing and retrieving a unique 

28# key in a thread local store. 

29_get_thread_key = threading.get_ident 

30 

31 

32# TODO(mdan): Move these to C++ as well. 

33# Moving to C++ can further avoid extra copies made by get_effective_map. 

34_source_mapper_stacks = collections.defaultdict(lambda: [SentinelMapper()]) 

35_source_filter_stacks = collections.defaultdict(lambda: [SentinelFilter()]) 

36 

37 

38class StackTraceTransform(object): 

39 """Base class for stack trace transformation functions.""" 

40 

41 _stack_dict = None # Subclasses should override 

42 _thread_key = None 

43 

44 def __enter__(self): 

45 # Any given instance is assumed to be used by a single thread, which reduces 

46 # expensive thread local lookups. 

47 if self._thread_key is None: 

48 self._thread_key = _get_thread_key() 

49 else: 

50 assert self._thread_key == _get_thread_key(), 'Shared across threads?' 

51 

52 stack = self._stack_dict[self._thread_key] 

53 self.parent = stack[-1] 

54 stack.append(self) 

55 self.update() 

56 return self 

57 

58 def __exit__(self, unused_type, unused_value, unused_traceback): 

59 top = self._stack_dict[self._thread_key].pop() 

60 assert top is self, 'Concurrent access?' 

61 

62 def update(self): 

63 raise NotImplementedError('subclasses need to override this') 

64 

65 

66class StackTraceMapper(StackTraceTransform): 

67 """Allows remapping traceback information to different source code.""" 

68 _stack_dict = _source_mapper_stacks 

69 

70 def __init__(self): 

71 self.internal_map = _tf_stack.PyBindSourceMap() 

72 

73 def update(self): 

74 self.internal_map.update_to(tuple(self.get_effective_source_map().items())) 

75 

76 def get_effective_source_map(self): 

77 """Returns a map (filename, lineno) -> (filename, lineno, function_name).""" 

78 raise NotImplementedError('subclasses need to override this') 

79 

80 

81EMPTY_DICT = {} 

82 

83 

84class SentinelMapper(StackTraceMapper): 

85 

86 def get_effective_source_map(self): 

87 return EMPTY_DICT 

88 

89 

90class StackTraceFilter(StackTraceTransform): 

91 """Allows filtering traceback information by removing superfluous frames.""" 

92 _stack_dict = _source_filter_stacks 

93 

94 def __init__(self): 

95 self.internal_set = _tf_stack.PyBindFileSet() 

96 

97 def update(self): 

98 self.internal_set.update_to(set(self.get_filtered_filenames())) 

99 

100 def get_filtered_filenames(self): 

101 raise NotImplementedError('subclasses need to override this') 

102 

103 

104EMPTY_SET = frozenset() 

105 

106 

107class SentinelFilter(StackTraceFilter): 

108 

109 def get_filtered_filenames(self): 

110 return EMPTY_SET 

111 

112 

113class CurrentModuleFilter(StackTraceFilter): 

114 """Filters stack frames from the module where this is used (best effort).""" 

115 

116 def __init__(self): 

117 super().__init__() 

118 filter_filename = None 

119 outer_f = None 

120 f = inspect.currentframe() 

121 try: 

122 if f is not None: 

123 # The current frame is __init__. The first outer frame should be the 

124 # caller. 

125 outer_f = f.f_back 

126 if outer_f is not None: 

127 filter_filename = inspect.getsourcefile(outer_f) 

128 self._filename = filter_filename 

129 # This may be called repeatedly: once on entry by the superclass, then by 

130 # each child context manager. 

131 self._cached_set = None 

132 finally: 

133 # Avoid reference cycles, see: 

134 # https://docs.python.org/3.7/library/inspect.html#the-interpreter-stack 

135 del f 

136 del outer_f 

137 

138 def get_filtered_filenames(self): 

139 if self._cached_set is not None: 

140 return self._cached_set 

141 

142 filtered_filenames = frozenset((self._filename,)) 

143 if self.parent is not None: 

144 filtered_filenames |= self.parent.get_filtered_filenames() 

145 self._cached_set = filtered_filenames 

146 return filtered_filenames 

147 

148 

149def extract_stack(stacklevel=1): 

150 """An eager-friendly alternative to traceback.extract_stack. 

151 

152 Args: 

153 stacklevel: number of initial frames to skip when producing the stack. 

154 

155 Returns: 

156 A list-like FrameSummary containing StackFrame-like objects, which are 

157 namedtuple-like objects with the following fields: filename, lineno, name, 

158 line, meant to masquerade as traceback.FrameSummary objects. 

159 """ 

160 thread_key = _get_thread_key() 

161 return _tf_stack.extract_stack( 

162 _source_mapper_stacks[thread_key][-1].internal_map, 

163 _source_filter_stacks[thread_key][-1].internal_set, 

164 stacklevel, 

165 ) 

166 

167 

168StackSummary = _tf_stack.StackTrace 

169FrameSummary = _tf_stack.StackFrame