From b73fea6e27b17dc4349754e585113e7a8138339e Mon Sep 17 00:00:00 2001
From: Tim Harley <tharley@google.com>
Date: Tue, 30 May 2017 08:26:18 -0700
Subject: [PATCH] Refactor `tf.Operation.traceback` implementation in to
methods of tf.Graph.
Adds an `_extract_frame_info` method to allow derived classes to extend the
information available in each op traceback, if desired. The default result of
`tf.Operation.traceback` is unchanged.
Also fixes a poorly scoped `pylint disable=line-too-long`, so adds the necessary
enable/disable blocks to silence pylint for the offending docstrings.
PiperOrigin-RevId: 157466174
---
tensorflow/python/framework/ops.py | 157 ++++++++++++++++-------------
1 file changed, 88 insertions(+), 69 deletions(-)
diff --git a/tensorflow/python/framework/ops.py b/tensorflow/python/framework/ops.py
index dfc3cf6dfd7..46417c23246 100644
--- a/tensorflow/python/framework/ops.py
+++ b/tensorflow/python/framework/ops.py
@@ -88,70 +88,6 @@ def _override_helper(clazz_object, operator, func):
setattr(clazz_object, operator, func)
-def _convert_stack(stack, include_func_start_lineno=False):
- """Converts a stack extracted using _extract_stack() to a traceback stack.
-
- Args:
- stack: A list of n 5-tuples,
- (filename, lineno, name, frame_globals, func_start_lineno).
- include_func_start_lineno: True if function start line number should be
- included as the 5th entry in return tuples.
-
- Returns:
- A list of n 4-tuples or 5-tuples
- (filename, lineno, name, code, [optional: func_start_lineno]), where the
- code tuple element is calculated from the corresponding elements of the
- input tuple.
- """
- ret = []
- for filename, lineno, name, frame_globals, func_start_lineno in stack:
- linecache.checkcache(filename)
- line = linecache.getline(filename, lineno, frame_globals)
- if line:
- line = line.strip()
- else:
- line = None
- if include_func_start_lineno:
- ret.append((filename, lineno, name, line, func_start_lineno))
- else:
- ret.append((filename, lineno, name, line))
- return ret
-
-
-# pylint: disable=line-too-long
-def _extract_stack():
- """A lightweight re-implementation of traceback.extract_stack.
-
- NOTE(mrry): traceback.extract_stack eagerly retrieves the line of code for
- each stack frame using linecache, which results in an abundance of stat()
- calls. This implementation does not retrieve the code, and any consumer
- should apply _convert_stack to the result to obtain a traceback that can
- be formatted etc. using traceback methods.
-
- Returns:
- A list of 5-tuples
- (filename, lineno, name, frame_globals, func_start_lineno) corresponding to
- the call stack of the current thread.
- """
- # pylint: enable=line-too-long
- try:
- raise ZeroDivisionError
- except ZeroDivisionError:
- f = sys.exc_info()[2].tb_frame.f_back
- ret = []
- while f is not None:
- lineno = f.f_lineno
- co = f.f_code
- filename = co.co_filename
- name = co.co_name
- frame_globals = f.f_globals
- func_start_lineno = co.co_firstlineno
- ret.append((filename, lineno, name, frame_globals, func_start_lineno))
- f = f.f_back
- ret.reverse()
- return ret
-
-
def _as_graph_element(obj):
"""Convert `obj` to a graph element if possible, otherwise return `None`.
@@ -1264,7 +1200,7 @@ class Operation(object):
self._original_op = original_op
self._op_def = op_def
- self._traceback = _extract_stack()
+ self._traceback = self._graph._extract_stack() # pylint: disable=protected-access
# Add this op to the current control flow context:
self._control_flow_context = g._get_control_flow_context()
if self._control_flow_context is not None:
@@ -1613,6 +1549,7 @@ class Operation(object):
@property
def node_def(self):
+ # pylint: disable=line-too-long
"""Returns a serialized `NodeDef` representation of this operation.
Returns:
@@ -1620,10 +1557,12 @@ class Operation(object):
[`NodeDef`](https://www.tensorflow.org/code/tensorflow/core/framework/node_def.proto)
protocol buffer.
"""
+ # pylint: enable=line-too-long
return self._node_def
@property
def op_def(self):
+ # pylint: disable=line-too-long
"""Returns the `OpDef` proto that represents the type of this op.
Returns:
@@ -1631,12 +1570,13 @@ class Operation(object):
[`OpDef`](https://www.tensorflow.org/code/tensorflow/core/framework/op_def.proto)
protocol buffer.
"""
+ # pylint: enable=line-too-long
return self._op_def
@property
def traceback(self):
"""Returns the call stack from when this operation was constructed."""
- return _convert_stack(self._traceback)
+ return self._graph._convert_stack(self._traceback) # pylint: disable=protected-access
@property
def traceback_with_start_lines(self):
@@ -1645,7 +1585,8 @@ class Operation(object):
Returns:
A list of 5-tuples (filename, lineno, name, code, func_start_lineno).
"""
- return _convert_stack(self._traceback, include_func_start_lineno=True)
+ return self._graph._convert_stack( # pylint: disable=protected-access
+ self._traceback, include_func_start_lineno=True)
def get_attr(self, name):
"""Returns the value of the attr of this op with the given `name`.
@@ -2170,6 +2111,76 @@ class Graph(object):
else:
self._scoped_c_graph = None
+ def _convert_stack(self, stack, include_func_start_lineno=False):
+ """Converts a stack extracted using _extract_stack() to a traceback stack.
+
+ Args:
+ stack: A list of n 5-tuples,
+ (filename, lineno, name, frame_globals, func_start_lineno).
+ include_func_start_lineno: True if function start line number should be
+ included as the 5th entry in return tuples.
+
+ Returns:
+ A list of n 4-tuples or 5-tuples
+ (filename, lineno, name, code, [optional: func_start_lineno]), where the
+ code tuple element is calculated from the corresponding elements of the
+ input tuple.
+ """
+ ret = []
+ for (filename, lineno, name, frame_globals, func_start_lineno,
+ unused_frame_info) in stack:
+ linecache.checkcache(filename)
+ line = linecache.getline(filename, lineno, frame_globals)
+ if line:
+ line = line.strip()
+ else:
+ line = None
+ if include_func_start_lineno:
+ ret.append((filename, lineno, name, line, func_start_lineno))
+ else:
+ ret.append((filename, lineno, name, line))
+ return ret
+
+ def _extract_stack(self):
+ """A lightweight, extensible re-implementation of traceback.extract_stack.
+
+ NOTE(mrry): traceback.extract_stack eagerly retrieves the line of code for
+ each stack frame using linecache, which results in an abundance of stat()
+ calls. This implementation does not retrieve the code, and any consumer
+ should apply _convert_stack to the result to obtain a traceback that can
+ be formatted etc. using traceback methods.
+
+ Derived classes can implement _extract_frame_info() to add extra information
+ to the traceback.
+
+ Returns:
+ A list of 6-tuples
+ (filename, lineno, name, frame_globals, func_start_lineno, custom_info)
+ corresponding to the call stack of the current thread.
+ """
+ try:
+ raise ZeroDivisionError
+ except ZeroDivisionError:
+ f = sys.exc_info()[2].tb_frame.f_back
+ ret = []
+ while f is not None:
+ lineno = f.f_lineno
+ co = f.f_code
+ filename = co.co_filename
+ name = co.co_name
+ frame_globals = f.f_globals
+ func_start_lineno = co.co_firstlineno
+ frame_info = self._extract_frame_info(f)
+ ret.append((filename, lineno, name, frame_globals, func_start_lineno,
+ frame_info))
+ f = f.f_back
+ ret.reverse()
+ return ret
+
+ def _extract_frame_info(self, frame): # pylint: disable=unused-argument
+ """Extracts custom information from a frame in an op traceback."""
+ return None
+
def _check_not_finalized(self):
"""Check if the graph is finalized.
@@ -2226,6 +2237,7 @@ class Graph(object):
@property
def graph_def_versions(self):
+ # pylint: disable=line-too-long
"""The GraphDef version information of this graph.
For details on the meaning of each version, see
@@ -2234,6 +2246,7 @@ class Graph(object):
Returns:
A `VersionDef`.
"""
+ # pylint: enable=line-too-long
return self._graph_def_versions
@property
@@ -2287,6 +2300,7 @@ class Graph(object):
self._control_flow_context = context
def _as_graph_def(self, from_version=None, add_shapes=False):
+ # pylint: disable=line-too-long
"""Returns a serialized `GraphDef` representation of this graph.
The serialized `GraphDef` can be imported into another `Graph`
@@ -2312,6 +2326,7 @@ class Graph(object):
ValueError: If the `graph_def` would be too large.
"""
+ # pylint: enable=line-too-long
with self._lock:
graph = graph_pb2.GraphDef()
graph.versions.CopyFrom(self._graph_def_versions)
@@ -2341,6 +2356,7 @@ class Graph(object):
return graph, self._version
def as_graph_def(self, from_version=None, add_shapes=False):
+ # pylint: disable=line-too-long
"""Returns a serialized `GraphDef` representation of this graph.
The serialized `GraphDef` can be imported into another `Graph`
@@ -2363,6 +2379,7 @@ class Graph(object):
Raises:
ValueError: If the `graph_def` would be too large.
"""
+ # pylint: enable=line-too-long
result, _ = self._as_graph_def(from_version, add_shapes)
return result
@@ -2930,7 +2947,7 @@ class Graph(object):
finally:
self._default_original_op = old_original_op
- # pylint: disable=g-doc-return-or-yield
+ # pylint: disable=g-doc-return-or-yield,line-too-long
@tf_contextlib.contextmanager
def name_scope(self, name):
r"""Returns a context manager that creates hierarchical names for operations.
@@ -3040,7 +3057,7 @@ class Graph(object):
yield "" if new_stack is None else new_stack + "/"
finally:
self._name_stack = old_stack
- # pylint: enable=g-doc-return-or-yield
+ # pylint: enable=g-doc-return-or-yield,line-too-long
def unique_name(self, name, mark_as_used=True):
"""Return a unique operation name for `name`.
@@ -3181,6 +3198,7 @@ class Graph(object):
@tf_contextlib.contextmanager
def device(self, device_name_or_function):
+ # pylint: disable=line-too-long
"""Returns a context manager that specifies the default device to use.
The `device_name_or_function` argument may either be a device name
@@ -3237,6 +3255,7 @@ class Graph(object):
created ops.
"""
+ # pylint: enable=line-too-long
if (device_name_or_function is not None
and not callable(device_name_or_function)):
device_function = pydev.merge_device(device_name_or_function)