[pypy-commit] pypy resource_warning: enable resource-tracking for socket objects

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: resource_warning
Changeset: r83568:777b0214d3a8
Date: 2016-04-07 16:08 +0200
http://bitbucket.org/pypy/pypy/changeset/777b0214d3a8/

Log:	enable resource-tracking for socket objects

diff --git a/pypy/interpreter/test/test_app_main.py b/pypy/interpreter/test/test_app_main.py
--- a/pypy/interpreter/test/test_app_main.py
+++ b/pypy/interpreter/test/test_app_main.py
@@ -1044,4 +1044,3 @@
             # assert it did not crash
         finally:
             sys.path[:] = old_sys_path
-    
diff --git a/pypy/module/_file/test/test_file.py b/pypy/module/_file/test/test_file.py
--- a/pypy/module/_file/test/test_file.py
+++ b/pypy/module/_file/test/test_file.py
@@ -11,6 +11,16 @@
             return file
     """)
 
+# the following function is used e.g. in test_resource_warning
+ at unwrap_spec(regex=str, s=str)
+def regex_search(space, regex, s):
+    import re
+    import textwrap
+    regex = textwrap.dedent(regex).strip()
+    m = re.search(regex, s)
+    m = bool(m)
+    return space.wrap(m)
+
 class AppTestFile(object):
     spaceconfig = dict(usemodules=("_file",))
 
@@ -18,16 +28,6 @@
         cls.w_temppath = cls.space.wrap(
             str(py.test.ensuretemp("fileimpl").join("foo.txt")))
         cls.w_file = getfile(cls.space)
-        #
-        # the following function is used e.g. in test_resource_warning
-        @unwrap_spec(regex=str, s=str)
-        def regex_search(space, regex, s):
-            import re
-            import textwrap
-            regex = textwrap.dedent(regex).strip()
-            m = re.search(regex, s)
-            m = bool(m)
-            return space.wrap(m)
         cls.w_regex_search = cls.space.wrap(interp2app(regex_search))
 
     def test_simple(self):
diff --git a/pypy/module/_socket/interp_socket.py b/pypy/module/_socket/interp_socket.py
--- a/pypy/module/_socket/interp_socket.py
+++ b/pypy/module/_socket/interp_socket.py
@@ -154,9 +154,28 @@
 
 
 class W_Socket(W_Root):
+    w_tb = None  # String representation of the traceback at creation time
+
     def __init__(self, space, sock):
+        self.space = space
         self.sock = sock
         register_socket(space, sock)
+        if self.space.sys.track_resources:
+            self.w_tb = self.space.format_traceback()
+
+    def __del__(self):
+        is_open = self.sock.fd >= 0
+        if is_open and self.space.sys.track_resources:
+            self.enqueue_for_destruction(self.space, W_Socket.destructor,
+                                         '__del__ method of ')
+
+    def destructor(self):
+        assert isinstance(self, W_Socket)
+        if self.space.sys.track_resources:
+            w_repr = self.space.repr(self)
+            str_repr = self.space.str_w(w_repr)
+            w_msg = self.space.wrap("WARNING: unclosed " + str_repr)
+            self.space.resource_warning(w_msg, self.w_tb)
 
     def get_type_w(self, space):
         return space.wrap(self.sock.type)
diff --git a/pypy/module/_socket/test/test_sock_app.py b/pypy/module/_socket/test/test_sock_app.py
--- a/pypy/module/_socket/test/test_sock_app.py
+++ b/pypy/module/_socket/test/test_sock_app.py
@@ -1,6 +1,8 @@
 import sys, os
 import py
 from pypy.tool.pytest.objspace import gettestobjspace
+from pypy.interpreter.gateway import interp2app
+from pypy.module._file.test.test_file import regex_search
 from rpython.tool.udir import udir
 from rpython.rlib import rsocket
 from rpython.rtyper.lltypesystem import lltype, rffi
@@ -314,6 +316,7 @@
     def setup_class(cls):
         cls.space = space
         cls.w_udir = space.wrap(str(udir))
+        cls.w_regex_search = space.wrap(interp2app(regex_search))
 
     def teardown_class(cls):
         if not cls.runappdirect:
@@ -402,6 +405,64 @@
         if os.name != 'nt':
             raises(OSError, os.close, fileno)
 
+    def test_socket_track_resources(self):
+        import _socket, os, gc, sys, cStringIO
+        s = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM, 0)
+        fileno = s.fileno()
+        assert s.fileno() >= 0
+        s.close()
+        assert s.fileno() < 0
+        s.close()
+        if os.name != 'nt':
+            raises(OSError, os.close, fileno)
+
+    def test_track_resources(self):
+        import os, gc, sys, cStringIO
+        import _socket
+        if '__pypy__' not in sys.builtin_module_names:
+            skip("pypy specific test")
+        #
+        def fn(flag1, flag2, do_close=False):
+            sys.pypy_set_track_resources(flag1)
+            mysock = _socket.socket(_socket.AF_INET, _socket.SOCK_STREAM, 0)
+            sys.pypy_set_track_resources(flag2)
+            buf = cStringIO.StringIO()
+            preverr = sys.stderr
+            try:
+                sys.stderr = buf
+                if do_close:
+                    mysock.close()
+                del mysock
+                gc.collect() # force __del__ to be called
+            finally:
+                sys.stderr = preverr
+                sys.pypy_set_track_resources(False)
+            return buf.getvalue()
+
+        # check with track_resources disabled
+        assert fn(False, False) == ""
+        #
+        # check that we don't get the warning if we actually closed the socket
+        msg = fn(True, True, do_close=True)
+        assert msg == ''
+        #
+        # check with track_resources enabled
+        msg = fn(True, True)
+        assert self.regex_search(r"""
+        WARNING: unclosed <socket object, .*>
+        Created at \(most recent call last\):
+          File ".*", line .*, in test_track_resources
+          File ".*", line .*, in fn
+          File ".*", line .*, in anonymous
+        """, msg)
+        #
+        # check with track_resources enabled in the destructor BUT with a
+        # file which was created when track_resources was disabled
+        msg = fn(False, True)
+        assert self.regex_search("WARNING: unclosed <socket object, .*>", msg)
+        assert "Created at" not in msg
+
+
     def test_socket_close_error(self):
         import _socket, os
         if os.name == 'nt':