Add reference leak check to Python well_known_types_test. Fix two refleak bugs.

-Fix ref leak when assign Struct field with creation
-Fix ref leak on "in" operator for ListValue

PiperOrigin-RevId: 759744611

Author: anandolee

python/google/protobuf/internal/well_known_types_test.py

@@ -16,6 +16,7 @@
 from google.protobuf import json_format
 from google.protobuf import text_format
 from google.protobuf.internal import more_messages_pb2
+from google.protobuf.internal import testing_refleaks
 from google.protobuf.internal import well_known_types
 from google.protobuf.internal import well_known_types_test_pb2
 
@@ -54,6 +55,7 @@ def CheckDurationConversion(self, message, text):
     self.assertEqual(message, parsed_message)
 
 
+@testing_refleaks.TestCase
 class TimeUtilTest(TimeUtilTestBase):
 
   def testTimestampSerializeAndParse(self):
@@ -704,6 +706,7 @@ def testDurationSub(self, old_time, time_delta, expected_sec, expected_nano):
     self.assertEqual(expected_nano, msg.optional_timestamp.nanos)
 
 
+@testing_refleaks.TestCase
 class StructTest(unittest.TestCase):
 
   def testEmptyDict(self):
@@ -987,6 +990,7 @@ def testMergeFrom(self):
     self.assertEqual(5, struct['key5'][0][1])
 
 
+@testing_refleaks.TestCase
 class AnyTest(unittest.TestCase):
 
   def testAnyMessage(self):

python/google/protobuf/pyext/message.cc

@@ -1124,14 +1124,16 @@ int InitAttributes(CMessage* self, PyObject* args, PyObject* kwargs) {
             Descriptor::WELLKNOWNTYPE_STRUCT) {
           ScopedPyObjectPtr ok(PyObject_CallMethod(
               reinterpret_cast<PyObject*>(cmessage), "update", "O", value));
-          if (ok.get() == nullptr && PyDict_Size(value) == 1 &&
-              PyDict_Contains(value, PyUnicode_FromString("fields"))) {
-            // Fallback to init as normal message field.
-            PyErr_Clear();
-            PyObject* tmp = Clear(cmessage);
-            Py_DECREF(tmp);
-            if (InitAttributes(cmessage, nullptr, value) < 0) {
-              return -1;
+          if (ok.get() == nullptr && PyDict_Size(value) == 1) {
+            ScopedPyObjectPtr fields_str(PyUnicode_FromString("fields"));
+            if (PyDict_Contains(value, fields_str.get())) {
+              // Fallback to init as normal message field.
+              PyErr_Clear();
+              PyObject* tmp = Clear(cmessage);
+              Py_DECREF(tmp);
+              if (InitAttributes(cmessage, nullptr, value) < 0) {
+                return -1;
+              }
             }
           }
         } else {
@@ -2391,21 +2393,19 @@ PyObject* Contains(CMessage* self, PyObject* arg) {
       const Reflection* reflection = message->GetReflection();
       const FieldDescriptor* map_field = descriptor->FindFieldByName("fields");
       const FieldDescriptor* key_field = map_field->message_type()->map_key();
-      PyObject* py_string = CheckString(arg, key_field);
-      if (!py_string) {
+      ScopedPyObjectPtr py_string(CheckString(arg, key_field));
+      if (py_string.get() == nullptr) {
         PyErr_SetString(PyExc_TypeError,
                         "The key passed to Struct message must be a str.");
         return nullptr;
       }
       char* value;
       Py_ssize_t value_len;
-      if (PyBytes_AsStringAndSize(py_string, &value, &value_len) < 0) {
-        Py_DECREF(py_string);
+      if (PyBytes_AsStringAndSize(py_string.get(), &value, &value_len) < 0) {
         Py_RETURN_FALSE;
       }
       std::string key_str;
       key_str.assign(value, value_len);
-      Py_DECREF(py_string);
 
       MapKey map_key;
       map_key.SetStringValue(key_str);
@@ -2414,9 +2414,9 @@ PyObject* Contains(CMessage* self, PyObject* arg) {
     }
     case Descriptor::WELLKNOWNTYPE_LISTVALUE: {
       // For WKT ListValue, check if the key is in the items.
-      PyObject* items = PyObject_CallMethod(reinterpret_cast<PyObject*>(self),
-                                            "items", nullptr);
-      return PyBool_FromLong(PySequence_Contains(items, arg));
+      ScopedPyObjectPtr items(PyObject_CallMethod(
+          reinterpret_cast<PyObject*>(self), "items", nullptr));
+      return PyBool_FromLong(PySequence_Contains(items.get(), arg));
     }
     default:
       // For other messages, check with HasField.

python/message.c

@@ -453,13 +453,16 @@ static bool PyUpb_Message_InitMessageAttribute(PyObject* _self, PyObject* name,
     const upb_MessageDef* msgdef = upb_FieldDef_MessageSubDef(field);
     if (upb_MessageDef_WellKnownType(msgdef) == kUpb_WellKnown_Struct) {
       ok = PyObject_CallMethod(submsg, "_internal_assign", "O", value);
-      if (!ok && PyDict_Size(value) == 1 &&
-          PyDict_Contains(value, PyUnicode_FromString("fields"))) {
-        // Fall back to init as normal message field.
-        PyErr_Clear();
-        PyObject* tmp = PyUpb_Message_Clear((PyUpb_Message*)submsg);
-        Py_DECREF(tmp);
-        ok = PyUpb_Message_InitAttributes(submsg, NULL, value) >= 0;
+      if (!ok && PyDict_Size(value) == 1) {
+        PyObject* fields_str = PyUnicode_FromString("fields");
+        if (PyDict_Contains(value, fields_str)) {
+          // Fall back to init as normal message field.
+          PyErr_Clear();
+          PyObject* tmp = PyUpb_Message_Clear((PyUpb_Message*)submsg);
+          Py_DECREF(tmp);
+          ok = PyUpb_Message_InitAttributes(submsg, NULL, value) >= 0;
+        }
+        Py_DECREF(fields_str);
       }
     } else {
       ok = PyUpb_Message_InitAttributes(submsg, NULL, value) >= 0;
@@ -1110,7 +1113,9 @@ static PyObject* PyUpb_Message_Contains(PyObject* _self, PyObject* arg) {
       PyUpb_Message* self = (void*)_self;
       if (PyUpb_Message_IsStub(self)) Py_RETURN_FALSE;
       PyObject* items = PyObject_CallMethod(_self, "items", NULL);
-      return PyBool_FromLong(PySequence_Contains(items, arg));
+      int ret = PySequence_Contains(items, arg);
+      Py_DECREF(items);
+      return PyBool_FromLong(ret);
     }
     default:
       // For other messages, check with HasField.