make _socket.socket weakrefable (closes #22569)

Patch from Alex Gaynor.

3 files changed, 15 insertions, 3 deletions

diff --git a/Lib/test/test_socket.py b/Lib/test/test_socket.py
index 81bd537..0ce80b9 100644
--- a/Lib/test/test_socket.py
+++ b/Lib/test/test_socket.py
@@ -12,9 +12,9 @@ import sys
 import os
 import array
 import contextlib
-from weakref import proxy
 import signal
 import math
+import weakref
 try:
     import _socket
 except ImportError:
@@ -264,7 +264,7 @@ class GeneralModuleTests(unittest.TestCase):
 
     def test_weakref(self):
         s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        p = proxy(s)
+        p = weakref.proxy(s)
         self.assertEqual(p.fileno(), s.fileno())
         s.close()
         s = None
@@ -275,6 +275,14 @@ class GeneralModuleTests(unittest.TestCase):
         else:
             self.fail('Socket proxy still exists')
 
+    def test_weakref__sock(self):
+        s = socket.socket()._sock
+        w = weakref.ref(s)
+        self.assertIs(w(), s)
+        del s
+        test_support.gc_collect()
+        self.assertIsNone(w())
+
     def testSocketError(self):
         # Testing socket module exceptions
         def raise_error(*args, **kwargs):
diff --git a/Modules/socketmodule.c b/Modules/socketmodule.c
index 0127a6c..e9e4479 100644
--- a/Modules/socketmodule.c
+++ b/Modules/socketmodule.c
@@ -3115,6 +3115,8 @@ sock_dealloc(PySocketSockObject *s)
 {
     if (s->sock_fd != -1)
         (void) SOCKETCLOSE(s->sock_fd);
+    if (s->weakreflist != NULL)
+        PyObject_ClearWeakRefs((PyObject *)s);
     Py_TYPE(s)->tp_free((PyObject *)s);
 }
 
@@ -3163,6 +3165,7 @@ sock_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
         ((PySocketSockObject *)new)->sock_fd = -1;
         ((PySocketSockObject *)new)->sock_timeout = -1.0;
         ((PySocketSockObject *)new)->errorhandler = &set_error;
+        ((PySocketSockObject *)new)->weakreflist = NULL;
     }
     return new;
 }
@@ -3226,7 +3229,7 @@ static PyTypeObject sock_type = {
     0,                                          /* tp_traverse */
     0,                                          /* tp_clear */
     0,                                          /* tp_richcompare */
-    0,                                          /* tp_weaklistoffset */
+    offsetof(PySocketSockObject, weakreflist),  /* tp_weaklistoffset */
     0,                                          /* tp_iter */
     0,                                          /* tp_iternext */
     sock_methods,                               /* tp_methods */
diff --git a/Modules/socketmodule.h b/Modules/socketmodule.h
index 8515499..d98e00e 100644
--- a/Modules/socketmodule.h
+++ b/Modules/socketmodule.h
@@ -132,6 +132,7 @@ typedef struct {
                                         sets a Python exception */
     double sock_timeout;                 /* Operation timeout in seconds;
                                         0.0 means non-blocking */
+    PyObject *weakreflist;
 } PySocketSockObject;
 
 /* --- C API ----------------------------------------------------*/