public:

PyBindSourceMap() : source_map_(std::make_shared<SourceMap>()) {}

// Shares ownership with whoever captures traces in the scope of this map.

std::shared_ptr<SourceMap> source_map_;

public:

PyBindFileSet() : file_set_(std::make_shared<StringSet>()) {}

// Shares ownership with whoever captures traces in the scope of this set.

std::shared_ptr<StringSet> file_set_;

public:

StackTraceWrapper(const std::shared_ptr<StackTrace>& captured,

const std::shared_ptr<SourceMap>& source_map,

const std::shared_ptr<StringSet>& filter, int stacklevel)

: captured_(captured),

source_map_(source_map),

filter_(filter),

stacklevel_(stacklevel) {}

~StackTraceWrapper() override {

PyGILState_STATE state = PyGILState_Ensure();

captured_.reset();

source_map_.reset();

filter_.reset();

PyGILState_Release(state);

}

StackTraceWrapper(StackTraceWrapper&& rhs) = default;

StackTraceWrapper& operator=(StackTraceWrapper&& rhs) = default;

static std::unique_ptr<StackTraceWrapper> ExtractStack(

const std::shared_ptr<SourceMap>& source_map,

const std::shared_ptr<StringSet>& filter, int stacklevel) {

return std::make_unique<StackTraceWrapper>(StackTrace::Capture(-1),

source_map, filter, stacklevel);

}

absl::Span<const StackFrame> ToFrames() const override {

ComputeFrozen();

return cache_->ToFrames();

}

std::vector<StackFrame> ToUncachedFrames() const override {

std::vector<StackFrame> frames = captured_->ToStackFrames(

*source_map_, [&](const char* f) { return StackTraceFiltering(f); },

/*reverse_traversal=*/false, /*limit=*/-1);

// Drop last stack frames.

int newsize = frames.size() - stacklevel_;

if (newsize < 0) {

newsize = 0;

}

frames.resize(newsize);

return frames;

}

std::vector<StackFrame> GetUserFrames(int limit) const override {

ComputeFrozen();

return cache_->GetUserFrames(limit);

}

StackFrame LastUserFrame() const override {

ComputeFrozen();

return cache_->LastUserFrame();

}

std::string ToString(const TracePrintingOptions& opts) const override {

ComputeFrozen();

return cache_->ToString(opts);

}

private:

void ComputeFrozen() const {

tsl::mutex_lock lock(mu_);

if (cache_ != nullptr) {

return;

}

std::vector<StackFrame> frames = ToUncachedFrames();

std::vector<StackFrame> user_frames = captured_->ToStackFrames(

*source_map_,

[&](const char* file_name) {

return StackTraceFiltering(file_name) ||

IsInternalFrameForFilename(file_name);

},

/*reverse_traversal=*/true,

/*limit=*/-1);

// ensure we use the original (outermost first) ordering.

absl::c_reverse(user_frames);

cache_ = std::make_unique<FrozenStackTrace>(frames, user_frames);

}

bool StackTraceFiltering(const char* file_name) const {

return filter_->contains(file_name);

}

mutable mutex mu_;

mutable std::unique_ptr<FrozenStackTrace> cache_;

std::shared_ptr<const StackTrace> captured_;

std::shared_ptr<SourceMap> source_map_;

std::shared_ptr<StringSet> filter_;

int stacklevel_;

pybind11::google::ImportStatusModule();

py::class_<PyBindSourceMap>(m, "PyBindSourceMap")

.def(py::init())

.def("update_to",

[](const PyBindSourceMap& self, const py::tuple& source_map) {

self.source_map_->clear();

for (const auto& item : source_map) {

const auto& tuple_item = py::cast<py::tuple>(item);

const auto& key = py::cast<py::tuple>(tuple_item[0]);

std::string&& k_filename = py::cast<std::string>(key[0]);

int k_lineno = py::cast<int>(key[1]);

const auto& value = py::cast<py::tuple>(tuple_item[1]);

std::string&& v_filename = py::cast<std::string>(value[0]);

int v_lineno = py::cast<int>(value[1]);

const auto& function_name_val = value[2];

std::string&& v_function_name =

function_name_val.is_none()

? ""

: py::cast<std::string>(function_name_val);

self.source_map_->emplace(

SourceLoc{k_filename, k_lineno},

StackFrame({v_filename, v_lineno, v_function_name}));

}

});

py::class_<PyBindFileSet>(m, "PyBindFileSet")

.def(py::init())

.def("update_to", [](const PyBindFileSet& self, const py::set& file_set) {

self.file_set_->clear();

for (const auto& item : file_set) {

self.file_set_->insert(py::cast<std::string>(item));

}

});

py::class_<GraphDebugInfoBuilder>(m, "GraphDebugInfoBuilder")

.def(py::init())

.def(

"AppendGraphDebugInfo",

[](GraphDebugInfoBuilder& self, std::string fn_name,

py::bytes debug_info_str) {

return self.AppendGraphDebugInfoStr(fn_name, debug_info_str);

},

py::arg("prefix"), py::arg("debug_info"))

.def(

"AccumulateStackTrace",

[](GraphDebugInfoBuilder& self, std::string function, std::string op,

const AbstractStackTrace& trace) {

std::string key = absl::StrCat(op, "@", function);

self.AccumulateStackTrace(

std::make_shared<FrozenStackTrace>(trace.ToFrames()), key);

},

py::arg("function"), py::arg("op"), py::arg("trace"))

.def("Build", [](GraphDebugInfoBuilder& self) -> py::bytes {

return py::bytes(self.ToGraphDebugInfoStr());

});

py::class_<StackFrame>(m, "StackFrame")

.def_property_readonly(

"filename",

[](const StackFrame& self) { return py::str(self.file_name); })

.def_property_readonly(

"lineno",

[](const StackFrame& self) { return py::int_(self.line_number); })

.def_property_readonly(

"name",

[](const StackFrame& self) { return py::str(self.function_name); })

.def_property_readonly("line", [](const StackFrame& self) { return ""; })

.def("__eq__", &StackFrame::operator==)

.def("__ne__", &StackFrame::operator!=)

.def("__hash__",

[](const StackFrame& self) {

return absl::Hash<std::tuple<std::string, int, std::string>>()(

std::make_tuple(self.file_name, self.line_number,

self.function_name));

})

.def("__getitem__",

[](const StackFrame& self, const py::object& index) -> py::object {

return py::make_tuple(

py::str(self.file_name), py::int_(self.line_number),

py::str(self.function_name), py::str(""))[index];

})

.def("__iter__",

[](const StackFrame& self) -> py::iterator {

return py::iter(py::make_tuple(

py::str(self.file_name), py::int_(self.line_number),

py::str(self.function_name), py::str("")));

})

.def("__repr__",

[](const StackFrame& self) -> py::str {

return absl::StrFormat("File \"%s\", line %d, in %s",

self.file_name, self.line_number,

py::str(self.function_name));

})

.def("__len__", [](const StackFrame&) { return 4; });

py::class_<AbstractStackTrace, std::shared_ptr<AbstractStackTrace>>(

m, "StackTrace")

.def(

"__getitem__",

[](const AbstractStackTrace& self, py::ssize_t index) -> StackFrame {

absl::Span<const StackFrame> frames = self.ToFrames();

const size_t eff_index =

index < 0 ? frames.size() + index : static_cast<size_t>(index);

if (eff_index >= frames.size()) {

throw py::index_error();

}

return frames[eff_index];

},

py::return_value_policy::take_ownership)

.def(

"__getitem__",

[](const AbstractStackTrace& self,

py::slice slice) -> std::shared_ptr<AbstractStackTrace> {

absl::Span<const StackFrame> frames = self.ToFrames();

py::ssize_t start, stop, step, slicelength;

if (!slice.compute(frames.size(), &start, &stop, &step,

&slicelength)) {

throw py::error_already_set();

}

if (step == 1) {

return std::make_shared<FrozenStackTrace>(

frames.subspan(start, slicelength));

}

std::vector<StackFrame> out;

out.reserve(slicelength);

// Python slices allow negative indexing.

for (int i = start; i != stop; i += step) {

out.push_back(frames[i]);

}

return std::make_shared<FrozenStackTrace>(out);

},

py::return_value_policy::take_ownership)

.def(

"__len__",

[](const AbstractStackTrace& self) { return self.ToFrames().size(); })

.def("__eq__",

[](const AbstractStackTrace& self, const AbstractStackTrace& other) {

return self.ToFrames() == other.ToFrames();

})

.def("__hash__",

[](const AbstractStackTrace& self) {

return py::hash(py::str(self.ToString({})));

})

.def(

"get_user_frames",

[](const AbstractStackTrace& self)

-> std::shared_ptr<AbstractStackTrace> {

return std::make_shared<FrozenStackTrace>(self.GetUserFrames(-1));

},

"Returns the non-framework frames as a new trace object.")

.def(

"last_user_frame",

[](const AbstractStackTrace& self) { return self.LastUserFrame(); },

"Returns the last non-framework frame.")

.def("__repr__",

[](const AbstractStackTrace& self) { return self.ToString({}); });

m.def(

"extract_stack",

[](const PyBindSourceMap& source_map, const PyBindFileSet& file_set,

int stacklevel) -> std::shared_ptr<AbstractStackTrace> {

return StackTraceWrapper::ExtractStack(source_map.source_map_,

file_set.file_set_, stacklevel);

},

py::arg("source_map"), py::arg("file_set"), py::arg("stacklevel") = 1,

py::return_value_policy::take_ownership);

m.def(

"LoadTracesFromDebugInfo",

[](py::bytes data) { return LoadTracesFromDebugInfoStr(data); },

py::arg("debug_info_proto"));