diff --git a/oneflow/api/python/functional/python_return_types.h b/oneflow/api/python/functional/python_return_types.h index c65efcb9084..df8de893109 100644 --- a/oneflow/api/python/functional/python_return_types.h +++ b/oneflow/api/python/functional/python_return_types.h @@ -32,10 +32,19 @@ namespace oneflow { namespace one { namespace functional { +inline PyObject* toTuple(PyStructSequence* obj) { +#if PY_MAJOR_VERSION == 2 + ROF_RUNTIME_ERROR() << "Oneflow do not support python 2"; +#else + Py_INCREF(obj); + return (PyObject*)obj; +#endif +} + PyObject* returned_structseq_repr(PyStructSequence* obj) { HANDLE_ERRORS PyTypeObject* tp = Py_TYPE(obj); - PyObjectPtr tuple((PyObject*)obj); + PyObject* tuple = toTuple(obj); if (tuple == nullptr) { return nullptr; } std::stringstream ss; @@ -49,23 +58,35 @@ PyObject* returned_structseq_repr(PyStructSequence* obj) { "In structseq_repr(), member %zd name is nullptr" " for type %.500s", i, tp->tp_name); + Py_DECREF(tuple); return nullptr; } - PyObject* val = PyTuple_GetItem(tuple.get(), i); - if (val == nullptr) { return nullptr; } + PyObject* val = PyTuple_GetItem(tuple, i); + if (val == nullptr) { + Py_DECREF(tuple); + return nullptr; + } - auto repr = PyObjectPtr(PyObject_Repr(val)); - if (repr == nullptr) { return nullptr; } + auto repr = PyObject_Repr(val); + if (repr == nullptr) { + Py_DECREF(tuple); + return nullptr; + } - const char* crepr = PyUnicode_AsUTF8(repr.get()); - if (crepr == nullptr) { return nullptr; } + const char* crepr = PyUnicode_AsUTF8(repr); + Py_DECREF(repr); + if (crepr == nullptr) { + Py_DECREF(tuple); + return nullptr; + } ss << cname << '=' << crepr; if (i < num_elements - 1) { ss << ",\n"; } } ss << ")"; + Py_DECREF(tuple); return PyUnicode_FromString(ss.str().c_str()); END_HANDLE_ERRORS }