| 41 | }; |
| 42 | |
| 43 | std::vector<StackFrame> ExtractStack(ssize_t limit, const py::list& mappers, |
| 44 | const py::list& filters) { |
| 45 | const py::dict& source_map = |
| 46 | mappers.size() == 0 |
| 47 | ? py::dict() |
| 48 | : mappers[mappers.size() - 1].attr("get_effective_source_map")(); |
| 49 | const py::set& filtered_filenames = |
| 50 | filters.size() == 0 |
| 51 | ? py::set() |
| 52 | : filters[filters.size() - 1].attr("get_filtered_filenames")(); |
| 53 | |
| 54 | const auto* tstate = PyThreadState_GET(); |
| 55 | // Drop extract_stack() wrapper-function frame from the result. |
| 56 | const PyFrameObject* f = tstate->frame->f_back; // TODO(slebedev): INCREF? |
| 57 | |
| 58 | std::vector<StackFrame> ret; |
| 59 | // 16 is somewhat arbitrary, but TensorFlow stack traces tend to be deep. |
| 60 | ret.reserve(limit < 0 ? 16 : static_cast<size_t>(limit)); |
| 61 | for (; f != nullptr && (limit < 0 || ret.size() < limit); f = f->f_back) { |
| 62 | PyCodeObject* co = f->f_code; |
| 63 | int lineno = PyFrame_GetLineNumber(const_cast<PyFrameObject*>(f)); |
| 64 | auto filename = py::reinterpret_borrow<py::str>(co->co_filename); |
| 65 | auto name = py::reinterpret_borrow<py::str>(co->co_name); |
| 66 | |
| 67 | // TODO(slebedev): consider moving the mappers/filters to C++ as well. |
| 68 | if (source_map.size() > 0) { |
| 69 | const auto& key = py::make_tuple(filename, lineno); |
| 70 | if (source_map.contains(key)) { |
| 71 | const py::tuple& mapped = source_map[key]; |
| 72 | filename = mapped[0]; |
| 73 | lineno = py::cast<py::int_>(mapped[1]); |
| 74 | name = mapped[2]; |
| 75 | } |
| 76 | } |
| 77 | |
| 78 | // Never filter the innermost frame. |
| 79 | // TODO(slebedev): upstream py::set::contains to pybind11. |
| 80 | if (!ret.empty() && |
| 81 | PySet_Contains(filtered_filenames.ptr(), filename.ptr())) |
| 82 | continue; |
| 83 | |
| 84 | const auto& globals = py::reinterpret_borrow<py::object>(f->f_globals); |
| 85 | const int func_start_lineno = co->co_firstlineno; |
| 86 | ret.push_back({std::move(filename), lineno, std::move(name), globals, |
| 87 | func_start_lineno}); |
| 88 | } |
| 89 | |
| 90 | std::reverse(ret.begin(), ret.end()); |
| 91 | return ret; |
| 92 | } |
| 93 | |
| 94 | } // namespace |
| 95 | |