Moved line lookup from tf_stack.convert_stack to StackFrame

This allows to align StackFrame interface with that of traceback.FrameSummary
in Python 3.5+.

PiperOrigin-RevId: 264705765

tensorflow/python/util/tf_stack.cc

@@ -38,6 +38,25 @@ struct StackFrame {
   py::str name;
   py::object globals;
   int func_start_lineno;
+
+  py::object line() const {
+    static const auto* linecache =
+        new py::module(py::module::import("linecache"));
+    const auto& checkcache = linecache->attr("checkcache");
+    const auto& getline = linecache->attr("getline");
+    checkcache(filename);
+    const auto& code =
+        py::cast<py::str>(getline(filename, lineno, globals).attr("strip")());
+    ssize_t size = 0;
+#if PY_MAJOR_VERSION == 3
+    if (PyUnicode_AsUTF8AndSize(code.ptr(), &size) == nullptr) {
+      throw py::error_already_set();
+    }
+#else
+    size = PyString_Size(code.ptr());
+#endif
+    return size > 0 ? static_cast<py::object>(code) : py::none();
+  }
 };
 
 std::vector<StackFrame> ExtractStack(ssize_t limit, const py::list& mappers,
@@ -59,7 +78,7 @@ std::vector<StackFrame> ExtractStack(ssize_t limit, const py::list& mappers,
   // 16 is somewhat arbitrary, but TensorFlow stack traces tend to be deep.
   ret.reserve(limit < 0 ? 16 : static_cast<size_t>(limit));
   for (; f != nullptr && (limit < 0 || ret.size() < limit); f = f->f_back) {
-    PyCodeObject* co = f->f_code;
+    const PyCodeObject* co = f->f_code;
     int lineno = PyFrame_GetLineNumber(const_cast<PyFrameObject*>(f));
     auto filename = py::reinterpret_borrow<py::str>(co->co_filename);
     auto name = py::reinterpret_borrow<py::str>(co->co_name);
@@ -75,11 +94,11 @@ std::vector<StackFrame> ExtractStack(ssize_t limit, const py::list& mappers,
       }
     }
 
-    // Never filter the innermost frame.
-    // TODO(slebedev): upstream py::set::contains to pybind11.
-    if (!ret.empty() &&
-        PySet_Contains(filtered_filenames.ptr(), filename.ptr()))
+    if (!ret.empty() &&  // Never filter the innermost frame.
+        filtered_filenames.size() > 0 &&
+        PySet_Contains(filtered_filenames.ptr(), filename.ptr())) {
       continue;
+    }
 
     const auto& globals = py::reinterpret_borrow<py::object>(f->f_globals);
     const int func_start_lineno = co->co_firstlineno;
@@ -94,24 +113,44 @@ std::vector<StackFrame> ExtractStack(ssize_t limit, const py::list& mappers,
 }  // namespace
 
 PYBIND11_MODULE(_tf_stack, m) {
-  // TODO(slebedev): consider dropping convert_stack in favor of
-  // a lazily initialized StackFrame.code property (using linecache).
+  // TODO(slebedev): rename to FrameSummary to match Python 3.5+.
   py::class_<StackFrame>(m, "StackFrame")
       .def(py::init<const py::str&, int, const py::str&, const py::object&,
                     int>())
       .def_readonly("filename", &StackFrame::filename)
       .def_readonly("lineno", &StackFrame::lineno)
       .def_readonly("name", &StackFrame::name)
+      // TODO(slebedev): remove globals and make the constructor private.
       .def_readonly("globals", &StackFrame::globals)
       .def_readonly("func_start_lineno", &StackFrame::func_start_lineno)
-      .def("__repr__", [](const StackFrame& self) {
-        return py::str(
-                   "StackFrame(filename={}, lineno={}, name={}, globals={}, "
-                   "func_start_lineno={})")
-            .format(self.filename, self.lineno, self.name, self.globals,
-                    self.func_start_lineno);
+      .def_property_readonly("line", &StackFrame::line)
+      .def("__repr__",
+           [](const StackFrame& self) {
+             return py::str("<StackFrame file {}, line {} in {}>")
+                 .format(self.filename, self.lineno, self.name);
+           })
+
+      // For compatibility with the traceback module.
+      .def("__getitem__",
+           [](const StackFrame& self, ssize_t index) -> py::object {
+             switch (index >= 0 ? index : 4 + index) {
+               case 0:
+                 return self.filename;
+               case 1:
+                 return py::cast(self.lineno);
+               case 2:
+                 return self.name;
+               case 3:
+                 return self.line();
+               default:
+                 throw py::index_error();
+             }
+           })
+      .def("__len__", [](const StackFrame&) {
+        return 4;  // For compatibility with the traceback module.
       });
 
+  // TODO(slebedev): rename to StackSummary to match Python 3.5+.
   py::bind_vector<std::vector<StackFrame>>(m, "Stack", py::module_local(true));
 
   m.def("extract_stack", [](const py::object& limit, const py::list& mappers,

tensorflow/python/util/tf_stack.py

@@ -20,7 +20,6 @@ from __future__ import print_function
 
 import collections
 import inspect
-import linecache
 import threading
 
 import six
@@ -170,8 +169,8 @@ 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).
+    stack: A sequence of StackFrame objects,
+      (filename, lineno, name, globals, func_start_lineno).
     include_func_start_lineno: True if function start line number should be
       included as the 5th entry in return tuples.
 
@@ -185,15 +184,11 @@ def convert_stack(stack, include_func_start_lineno=False):
     for frame in stack:
       filename = frame.filename
       lineno = frame.lineno
-      linecache.checkcache(filename)
-      line = linecache.getline(filename, lineno, frame.globals)
-      if line:
-        line = line.strip()
-      else:
-        line = None
+      name = frame.name
+      line = frame.line
       if include_func_start_lineno:
-        yield (filename, lineno, frame.name, line, frame.func_start_lineno)
+        yield (filename, lineno, name, line, frame.func_start_lineno)
       else:
-        yield (filename, lineno, frame.name, line)
+        yield (filename, lineno, name, line)
 
   return tuple(_tuple_generator())