Change TFStackTest to be Python3 compliant.

From Python3.5, we have a different object to represent a FrameSummary: https://docs.python.org/3/library/traceback.html#framesummary-objects

PiperOrigin-RevId: 276853881
Change-Id: I7563f0e0b14e71a76ad5554f48146d2c7221be65
This commit is contained in:
Mihai Maruseac 2019-10-26 06:34:22 -07:00 committed by TensorFlower Gardener
parent 3a33d697dd
commit c6920fcafc

View File

@ -36,7 +36,7 @@ class TFStackTest(test.TestCase):
def testConsistencyWithTraceback(self):
stack, expected_stack = extract_stack()
for frame, expected in zip(stack, expected_stack):
self.assertEqual(tuple(frame), expected)
self.assertEqual(convert_stack_frame(frame), expected)
def testFormatStack(self):
stack, expected_stack = extract_stack()
@ -58,5 +58,17 @@ def extract_stack(limit=None):
return tf_stack.extract_stack(limit), traceback.extract_stack(limit)
def convert_stack_frame(frame):
"""Converts a TF stack frame into Python's."""
# TODO(mihaimaruseac): Remove except case when dropping suport for py2
try:
return traceback.FrameSummary(
frame.filename, frame.lineno, frame.name, line=frame.line)
except AttributeError:
# On Python < 3.5 (i.e., Python2), we don't have traceback.FrameSummary so
# we don't need to match with that class. Instead, just a tuple is enough.
return tuple(frame)
if __name__ == "__main__":
test.main()