Allow passing a Request instance for the url parameter

Author: ephphatha

Lib/test/test_robotparser.py

@@ -773,6 +773,65 @@ def testServiceUnavailable(self):
         self.assertFalse(parser.can_fetch("*", url + '/path/file.html'))
 
 
+class UserAgentRobotHandler(BaseHTTPRequestHandler):
+
+    def do_GET(self):
+        if self.headers.get('User-Agent').startswith('Python-urllib'):
+            self.send_error(403, "Forbidden access")
+        else:
+            self.send_response(200)
+            self.end_headers()
+            self.wfile.write(b"User-agent: *\nDisallow:")
+
+    def log_message(self, format, *args):
+        pass
+
+
+@unittest.skipUnless(
+    support.has_socket_support,
+    "Socket server requires working socket."
+)
+class UserAgentSiteTestCase(unittest.TestCase):
+
+    def setUp(self):
+        # clear _opener global variable
+        self.addCleanup(urllib.request.urlcleanup)
+
+        self.server = HTTPServer((socket_helper.HOST, 0), UserAgentRobotHandler)
+
+        self.t = threading.Thread(
+            name='HTTPServer serving',
+            target=self.server.serve_forever,
+            # Short poll interval to make the test finish quickly.
+            # Time between requests is short enough that we won't wake
+            # up spuriously too many times.
+            kwargs={'poll_interval':0.01})
+        self.t.daemon = True  # In case this function raises.
+        self.t.start()
+
+    def tearDown(self):
+        self.server.shutdown()
+        self.t.join()
+        self.server.server_close()
+
+    @threading_helper.reap_threads
+    def testUserAgentFilteringSite(self):
+        addr = self.server.server_address
+        url = 'http://' + socket_helper.HOST + ':' + str(addr[1])
+        robots_url = url + "/robots.txt"
+        file_url = url + "/document"
+        parser = urllib.robotparser.RobotFileParser()
+        parser.set_url(robots_url)
+        parser.read()
+        self.assertTrue(parser.disallow_all)
+        self.assertFalse(parser.can_fetch("*", file_url))
+        parser = urllib.robotparser.RobotFileParser()
+        parser.set_url(urllib.request.Request(robots_url, headers={'User-Agent': 'cybermapper'}))
+        parser.read()
+        self.assertFalse(parser.disallow_all)
+        self.assertTrue(parser.can_fetch("*", file_url))
+
+
 @support.requires_working_socket()
 class NetworkTestCase(unittest.TestCase):
 

Lib/urllib/robotparser.py

@@ -57,7 +57,12 @@ def modified(self):
     def set_url(self, url):
         """Sets the URL referring to a robots.txt file."""
         self.url = url
-        self.host, self.path = urllib.parse.urlsplit(url)[1:3]
+
+        if isinstance(url, urllib.request.Request):
+            self.host = url.host
+            self.path = url.selector
+        else:
+            self.host, self.path = urllib.parse.urlparse(url)[1:3]
 
     def read(self):
         """Reads the robots.txt URL and feeds it to the parser."""