#1178141: add addinfourl.code to get http status code from urllib.

3 files changed, 24 insertions, 5 deletions

diff --git a/Lib/test/test_urllib.py b/Lib/test/test_urllib.py
index ed7f6ce..455e7fc 100644
--- a/Lib/test/test_urllib.py
+++ b/Lib/test/test_urllib.py
@@ -47,7 +47,7 @@ class urlopen_FileTests(unittest.TestCase):
     def test_interface(self):
         # Make sure object returned by urlopen() has the specified methods
         for attr in ("read", "readline", "readlines", "fileno",
-                     "close", "info", "geturl", "__iter__"):
+                     "close", "info", "geturl", "getcode", "__iter__"):
             self.assert_(hasattr(self.returned_obj, attr),
                          "object returned by urlopen() lacks %s attribute" %
                          attr)
@@ -87,6 +87,9 @@ class urlopen_FileTests(unittest.TestCase):
     def test_geturl(self):
         self.assertEqual(self.returned_obj.geturl(), self.pathname)
 
+    def test_getcode(self):
+        self.assertEqual(self.returned_obj.getcode(), None)
+
     def test_iter(self):
         # Test iterator
         # Don't need to count number of iterations since test would fail the
@@ -123,6 +126,8 @@ class urlopen_HttpTests(unittest.TestCase):
             fp = urllib.urlopen("http://python.org/")
             self.assertEqual(fp.readline(), 'Hello!')
             self.assertEqual(fp.readline(), '')
+            self.assertEqual(fp.geturl(), 'http://python.org/')
+            self.assertEqual(fp.getcode(), 200)
         finally:
             self.unfakehttp()
 
diff --git a/Lib/test/test_urllibnet.py b/Lib/test/test_urllibnet.py
index 9105afe..a28c081 100644
--- a/Lib/test/test_urllibnet.py
+++ b/Lib/test/test_urllibnet.py
@@ -83,6 +83,16 @@ class urlopenNetworkTests(unittest.TestCase):
             open_url.close()
         self.assertEqual(gotten_url, URL)
 
+    def test_getcode(self):
+        # test getcode() with the fancy opener to get 404 error codes
+        URL = "http://www.python.org/XXXinvalidXXX"
+        open_url = urllib.FancyURLopener().open(URL)
+        try:
+            code = open_url.getcode()
+        finally:
+            open_url.close()
+        self.assertEqual(code, 404)
+
     def test_fileno(self):
         if (sys.platform in ('win32',) or
                 not hasattr(os, 'fdopen')):
diff --git a/Lib/urllib.py b/Lib/urllib.py
index c6751b8..4a49fc0 100644
--- a/Lib/urllib.py
+++ b/Lib/urllib.py
@@ -343,7 +343,7 @@ class URLopener:
         # According to RFC 2616, "2xx" code indicates that the client's
         # request was successfully received, understood, and accepted.
         if (200 <= errcode < 300):
-            return addinfourl(fp, headers, "http:" + url)
+            return addinfourl(fp, headers, "http:" + url, errcode)
         else:
             if data is None:
                 return self.http_error(url, fp, errcode, errmsg, headers)
@@ -438,7 +438,7 @@ class URLopener:
             # According to RFC 2616, "2xx" code indicates that the client's
             # request was successfully received, understood, and accepted.
             if (200 <= errcode < 300):
-                return addinfourl(fp, headers, "https:" + url)
+                return addinfourl(fp, headers, "https:" + url, errcode)
             else:
                 if data is None:
                     return self.http_error(url, fp, errcode, errmsg, headers)
@@ -610,7 +610,7 @@ class FancyURLopener(URLopener):
 
     def http_error_default(self, url, fp, errcode, errmsg, headers):
         """Default error handling -- don't raise an exception."""
-        return addinfourl(fp, headers, "http:" + url)
+        return addinfourl(fp, headers, "http:" + url, errcode)
 
     def http_error_302(self, url, fp, errcode, errmsg, headers, data=None):
         """Error 302 -- relocated (temporarily)."""
@@ -953,14 +953,18 @@ class addinfo(addbase):
 class addinfourl(addbase):
     """class to add info() and geturl() methods to an open file."""
 
-    def __init__(self, fp, headers, url):
+    def __init__(self, fp, headers, url, code=None):
         addbase.__init__(self, fp)
         self.headers = headers
         self.url = url
+        self.code = code
 
     def info(self):
         return self.headers
 
+    def getcode(self):
+        return self.code
+
     def geturl(self):
         return self.url