urllib.response object to use _TemporaryFileWrapper (and _TemporaryFileCloser)

facility. Provides a better way to handle file descriptor close.

Address issue #15002 . Patch contributed by Christian Theune.

2 files changed, 60 insertions, 66 deletions

diff --git a/Lib/test/test_urllib_response.py b/Lib/test/test_urllib_response.py
index fdd3325..0eb5942 100644
--- a/Lib/test/test_urllib_response.py
+++ b/Lib/test/test_urllib_response.py
@@ -1,42 +1,59 @@
 """Unit tests for code in urllib.response."""
 
-import test.support
+import socket
+import tempfile
 import urllib.response
 import unittest
 
-class TestFile(object):
-
-    def __init__(self):
-        self.closed = False
-
-    def read(self, bytes):
-        pass
-
-    def readline(self):
-        pass
-
-    def close(self):
-        self.closed = True
-
-class Testaddbase(unittest.TestCase):
-
-    # TODO(jhylton): Write tests for other functionality of addbase()
+class TestResponse(unittest.TestCase):
 
     def setUp(self):
-        self.fp = TestFile()
-        self.addbase = urllib.response.addbase(self.fp)
+        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        self.fp = self.sock.makefile('rb')
+        self.test_headers = {"Host": "www.python.org",
+                             "Connection": "close"}
 
     def test_with(self):
+        addbase = urllib.response.addbase(self.fp)
+
+        self.assertIsInstance(addbase, tempfile._TemporaryFileWrapper)
+
         def f():
-            with self.addbase as spam:
+            with addbase as spam:
                 pass
         self.assertFalse(self.fp.closed)
         f()
         self.assertTrue(self.fp.closed)
         self.assertRaises(ValueError, f)
 
-def test_main():
-    test.support.run_unittest(Testaddbase)
+    def test_addclosehook(self):
+        closehook_called = False
+
+        def closehook():
+            nonlocal closehook_called
+            closehook_called = True
+
+        closehook = urllib.response.addclosehook(self.fp, closehook)
+        closehook.close()
+
+        self.assertTrue(self.fp.closed)
+        self.assertTrue(closehook_called)
+
+    def test_addinfo(self):
+        info = urllib.response.addinfo(self.fp, self.test_headers)
+        self.assertEqual(info.info(), self.test_headers)
+
+    def test_addinfourl(self):
+        url = "http://www.python.org"
+        code = 200
+        infourl = urllib.response.addinfourl(self.fp, self.test_headers,
+                                             url, code)
+        self.assertEqual(infourl.info(), self.test_headers)
+        self.assertEqual(infourl.geturl(), url)
+        self.assertEqual(infourl.getcode(), code)
+
+    def tearDown(self):
+        self.sock.close()
 
 if __name__ == '__main__':
-    test_main()
+    unittest.main()
diff --git a/Lib/urllib/response.py b/Lib/urllib/response.py
index 1cf1d1a..4a143b0 100644
--- a/Lib/urllib/response.py
+++ b/Lib/urllib/response.py
@@ -6,60 +6,39 @@ addinfourl instance, which defines an info() method that returns
 headers and a geturl() method that returns the url.
 """
 
-class addbase(object):
-    """Base class for addinfo and addclosehook."""
+import tempfile
+
+__all__ = ['addbase', 'addclosehook', 'addinfo', 'addinfourl']
+
+
+class addbase(tempfile._TemporaryFileWrapper):
+    """Base class for addinfo and addclosehook. Is a good idea for garbage collection."""
 
     # XXX Add a method to expose the timeout on the underlying socket?
 
     def __init__(self, fp):
-        # TODO(jhylton): Is there a better way to delegate using io?
+        super(addbase,  self).__init__(fp, '<urllib response>', delete=False)
+        # Keep reference around as this was part of the original API.
         self.fp = fp
-        self.read = self.fp.read
-        self.readline = self.fp.readline
-        # TODO(jhylton): Make sure an object with readlines() is also iterable
-        if hasattr(self.fp, "readlines"):
-            self.readlines = self.fp.readlines
-        if hasattr(self.fp, "fileno"):
-            self.fileno = self.fp.fileno
-        else:
-            self.fileno = lambda: None
-
-    def __iter__(self):
-        # Assigning `__iter__` to the instance doesn't work as intended
-        # because the iter builtin does something like `cls.__iter__(obj)`
-        # and thus fails to find the _bound_ method `obj.__iter__`.
-        # Returning just `self.fp` works for built-in file objects but
-        # might not work for general file-like objects.
-        return iter(self.fp)
 
     def __repr__(self):
         return '<%s at %r whose fp = %r>' % (self.__class__.__name__,
-                                             id(self), self.fp)
-
-    def close(self):
-        if self.fp:
-            self.fp.close()
-        self.fp = None
-        self.read = None
-        self.readline = None
-        self.readlines = None
-        self.fileno = None
-        self.__iter__ = None
-        self.__next__ = None
+                                             id(self), self.file)
 
     def __enter__(self):
-        if self.fp is None:
+        if self.fp.closed:
             raise ValueError("I/O operation on closed file")
         return self
 
     def __exit__(self, type, value, traceback):
         self.close()
 
+
 class addclosehook(addbase):
     """Class to add a close hook to an open file."""
 
     def __init__(self, fp, closehook, *hookargs):
-        addbase.__init__(self, fp)
+        super(addclosehook, self).__init__(fp)
         self.closehook = closehook
         self.hookargs = hookargs
 
@@ -68,30 +47,28 @@ class addclosehook(addbase):
             self.closehook(*self.hookargs)
             self.closehook = None
             self.hookargs = None
-        addbase.close(self)
+        super(addclosehook, self).close()
+
 
 class addinfo(addbase):
     """class to add an info() method to an open file."""
 
     def __init__(self, fp, headers):
-        addbase.__init__(self, fp)
+        super(addinfo, self).__init__(fp)
         self.headers = headers
 
     def info(self):
         return self.headers
 
-class addinfourl(addbase):
+
+class addinfourl(addinfo):
     """class to add info() and geturl() methods to an open file."""
 
     def __init__(self, fp, headers, url, code=None):
-        addbase.__init__(self, fp)
-        self.headers = headers
+        super(addinfourl, self).__init__(fp, headers)
         self.url = url
         self.code = code
 
-    def info(self):
-        return self.headers
-
     def getcode(self):
         return self.code