pub_soft/alphafold3
3
0
Fork 0
mirror of https://github.com/google-deepmind/alphafold3.git synced 2026-06-02 11:54:36 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix KeyError in string_array.remap_multiple for Python 3.14

Browse Source 
PiperOrigin-RevId: 914331813
Change-Id: I66c4693c285489e2ba8ad833ebb6e69b677cdb1d
This commit is contained in:
Augustin Zidek
2026-05-12 09:47:20 -07:00
committed by Copybara-Service
parent 8d283f3baf
commit 9ea7db2d3b
1 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
src/alphafold3/structure/cpp/string_array_pybind.cc
View File

@@ -174,6 +174,9 @@ py::array_t<bool> IsIn(
py::array RemapMultipleArrays(
const std::vector<py::array_t<PyObject*, py::array::c_style>>& arrays,
const py::dict& mapping) {
if (arrays.empty()) {
return py::array_t<int64_t>(0);
}
size_t array_size = arrays[0].size();
for (const auto& array : arrays) {
if (array.size() != array_size) {
@@ -184,10 +187,6 @@ py::array RemapMultipleArrays(
// Create a result buffer.
auto result = py::array_t<int64_t>(array_size);
absl::Span<int64_t> result_buffer(result.mutable_data(), array_size);
PyObject* entry = PyTuple_New(arrays.size());
if (entry == nullptr) {
throw py::error_already_set();
}
std::vector<absl::Span<PyObject* const>> array_spans;
array_spans.reserve(arrays.size());
for (const auto& array : arrays) {
@@ -197,34 +196,38 @@ py::array RemapMultipleArrays(
// Iterate over arrays and look up elements in the `py_dict`.
bool fail = false;
for (size_t i = 0; i < array_size; ++i) {
PyObject* entry = PyTuple_New(arrays.size());
if (entry == nullptr) {
fail = true;
break;
}
for (size_t j = 0; j < array_spans.size(); ++j) {
PyTuple_SET_ITEM(entry, j, array_spans[j][i]);
PyTuple_SET_ITEM(entry, j, Py_NewRef(array_spans[j][i]));
}
PyObject* result = PyDict_GetItem(mapping.ptr(), entry);
if (result != nullptr) {
int64_t result_value = PyLong_AsLongLong(result);
if (result_value == -1 && PyErr_Occurred()) {
fail = true;
Py_DECREF(entry);
break;
}
if (result_value > std::numeric_limits<int64_t>::max() ||
result_value < std::numeric_limits<int64_t>::lowest()) {
PyErr_SetString(PyExc_OverflowError, "Result value too large.");
fail = true;
Py_DECREF(entry);
break;
}
result_buffer[i] = result_value;
} else {
PyErr_Format(PyExc_KeyError, "%R", entry);
fail = true;
Py_DECREF(entry);
break;
}
Py_DECREF(entry);
}
for (size_t j = 0; j < array_spans.size(); ++j) {
PyTuple_SET_ITEM(entry, j, nullptr);
}
Py_XDECREF(entry);
if (fail) {
throw py::error_already_set();
}