Fix #359: request_headers raises error when grab instance did zero re…

…quests

grab/base.py

@@ -16,6 +16,7 @@
 from datetime import datetime
 import weakref
 from typing import Dict, Any, Optional, cast
+from email.message import EmailMessage
 
 from six.moves.urllib.parse import urljoin
 import six
@@ -810,11 +811,17 @@ def __setstate__(self, state):
             setattr(self, slot, value)
 
     @property
-    def request_headers(self):
-        first_head = self.request_head.decode("utf-8").split("\r\n\r\n")[0]
-        lines = first_head.split("\r\n")
-        lines = [x for x in lines if ":" in x]
-        return email.message_from_string("\n".join(lines))
+    def request_headers(self) -> Optional[EmailMessage]:
+        if self.request_head is None:
+            return None
+        else:
+            first_head = self.request_head.decode("utf-8").split("\r\n\r\n")[0]
+            lines = first_head.split("\r\n")
+            lines = [x for x in lines if ":" in x]
+            return cast(
+                EmailMessage,
+                email.message_from_string("\n".join(lines), _class=EmailMessage),
+            )
 
 
 # For backward compatibility

runtest.py

@@ -30,6 +30,7 @@
     "tests.grab_get_request",
     "tests.grab_post_request",
     "tests.grab_request",  # FIXME: fix tests excluded for urllib3
+    "tests.grab_request_headers",  # FIXME: fix tests excluded for urllib3
     "tests.grab_user_agent",
     "tests.grab_cookies",  # FIXME: fix tests excluded for urllib3
     "tests.grab_url_processing",

tests/grab_request_headers.py

@@ -0,0 +1,53 @@
+from pprint import pprint  # pylint: disable=unused-import
+from email.message import EmailMessage
+
+from test_server import Response  # pylint: disable=unused-import
+
+from tests.util import build_grab  # pylint: disable=unused-import
+from tests.util import BaseGrabTestCase
+
+
+class TestMisc(BaseGrabTestCase):
+    def setUp(self):
+        self.server.reset()
+
+    # *****************
+    # grab.request_head
+    # *****************
+
+    def test_request_headers_default(self):
+        grab = build_grab()
+        self.assertTrue(grab.request_headers is None)
+
+    def test_request_headers_debug_false(self):
+        self.server.add_response(Response())
+        grab = build_grab(debug=False)
+        grab.go(self.server.get_url(), headers=[("User-Agent", "GRAB")])
+        self.assertTrue(isinstance(grab.request_headers, EmailMessage))
+        self.assertTrue(len(grab.request_headers) == 0)
+
+    def test_request_headers_debug_true(self):
+        self.server.add_response(Response())
+        grab = build_grab(debug=True)
+        grab.go(self.server.get_url(), headers=[("User-Agent", "GRAB")])
+        self.assertEqual(grab.request_headers["user-agent"], "GRAB")
+
+    # *****************
+    # grab.request_head
+    # *****************
+
+    def test_request_head_default(self):
+        grab = build_grab()
+        self.assertTrue(grab.request_headers is None)
+
+    def test_request_head_debug_false(self):
+        self.server.add_response(Response())
+        grab = build_grab(debug=False)
+        grab.go(self.server.get_url())
+        self.assertEqual(grab.request_head, b"")
+
+    def test_request_head_debug_true(self):
+        self.server.add_response(Response())
+        grab = build_grab(debug=True)
+        grab.go(self.server.get_url())
+        self.assertTrue(grab.request_head.startswith(b"GET /"))

tests/misc.py

@@ -9,33 +9,3 @@
 class TestMisc(BaseGrabTestCase):
     def setUp(self):
         self.server.reset()
-
-    def test_put(self):
-        pass
-
-    # @only_grab_transport("pycurl")
-    # def test_different_domains(self):
-    #    import pycurl  # pylint: disable=import-outside-toplevel
-
-    #    grab = build_grab()
-    #    names = [
-    #        "foo:%d:127.0.0.1" % self.server.port,
-    #        "bar:%d:127.0.0.1" % self.server.port,
-    #    ]
-    #    grab.setup_transport("pycurl")
-    #    grab.transport.curl.setopt(pycurl.RESOLVE, names)
-
-    #    self.server.add_response(Response(headers=[("Set-Cookie", "foo=foo")]))
-    #    grab.go("http://foo:%d" % self.server.port)
-    #    self.assertEqual(dict(grab.doc.cookies.items()), {"foo": "foo"})
-
-    #    self.server.add_response(Response(headers=[("Set-Cookie", "bar=bar")]))
-    #    grab.go("http://bar:%d" % self.server.port)
-    #    self.assertEqual(dict(grab.doc.cookies.items()), {"bar": "bar"})
-
-    #    # That does not hold anymore, I guess I have fixed it
-    #    # # response.cookies contains cookies from both domains
-    #    # # because it just accumulates cookies over time
-    #    # # self.assertEqual(
-    #    # #     dict(grab.doc.cookies.items()), {"foo": "foo", "bar": "bar"}
-    #    # # )