http2: improve framereader

2 files changed, 10 insertions, 5 deletions

diff --git a/netlib/http/http2/__init__.py b/netlib/http/http2/__init__.py
index 60064190..7f84a1ab 100644
--- a/netlib/http/http2/__init__.py
+++ b/netlib/http/http2/__init__.py
@@ -1,8 +1,10 @@
 from __future__ import absolute_import, print_function, division
-from netlib.http.http2 import framereader
+
+from netlib.http.http2.framereader import read_raw_frame, parse_frame
 from netlib.http.http2.utils import parse_headers
 
 __all__ = [
-    "framereader",
+    "read_raw_frame",
+    "parse_frame",
     "parse_headers",
 ]
diff --git a/netlib/http/http2/framereader.py b/netlib/http/http2/framereader.py
index eb9b069a..8b7cfffb 100644
--- a/netlib/http/http2/framereader.py
+++ b/netlib/http/http2/framereader.py
@@ -4,7 +4,7 @@ import hyperframe
 from ...exceptions import HttpException
 
 
-def http2_read_raw_frame(rfile):
+def read_raw_frame(rfile):
     header = rfile.safe_read(9)
     length = int(codecs.encode(header[:3], 'hex_codec'), 16)
 
@@ -15,8 +15,11 @@ def http2_read_raw_frame(rfile):
     return [header, body]
 
 
-def http2_read_frame(rfile):
-    header, body = http2_read_raw_frame(rfile)
+def parse_frame(header, body=None):
+    if body is None:
+        body = header[9:]
+        header = header[:9]
+
     frame, length = hyperframe.frame.Frame.parse_frame_header(header)
     frame.parse_body(memoryview(body))
     return frame