fix_structseq_repr_error

Oneflow-Inc · May 20, 2024 · bb3fb19 · bb3fb19
1 parent ea585f6
commit bb3fb19
Showing 1 changed file with 28 additions and 7 deletions.
diff --git 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
 }