Added flag reading method for upgrade, to allow early discovery of an upgrade request (pre-raise, in on_message_complete)

Author: trollfot

httptools/parser/parser.pyx

@@ -151,6 +151,9 @@ cdef class HttpParser:
     def should_keep_alive(self):
         return bool(cparser.http_should_keep_alive(self._cparser))
 
+    def should_upgrade(self):
+        return bool(self._cparser.upgrade)
+
     def feed_data(self, data):
         cdef:
             size_t data_len

tests/test_parser.py

@@ -348,6 +348,26 @@ def test_parser_request_upgrade_1(self):
             b'Host': b'example.com',
             b'Upgrade': b'WebSocket'})
 
+    def test_parser_request_upgrade_flag(self):
+        m = mock.Mock()
+        p = httptools.HttpRequestParser(m)
+
+        def on_headers_complete():
+            self.assertEqual(p.should_upgrade(), False)
+
+        def on_message_complete():
+            self.assertEqual(p.should_upgrade(), True)
+
+        m.on_headers_complete = on_headers_complete
+        m.on_message_complete = on_message_complete
+    
+        try:
+            p.feed_data(UPGRADE_REQUEST1)
+        except httptools.HttpParserUpgrade as ex:
+            offset = ex.args[0]
+        else:
+            self.fail('HttpParserUpgrade was not raised')
+        
     def test_parser_request_error_in_on_header(self):
         class Error(Exception):
             pass