json contentview: minor improvements

2 files changed, 61 insertions, 36 deletions

diff --git a/mitmproxy/contentviews/json.py b/mitmproxy/contentviews/json.py
index dbe600aa..a5a7d03b 100644
--- a/mitmproxy/contentviews/json.py
+++ b/mitmproxy/contentviews/json.py
@@ -1,16 +1,39 @@
 import re
 import json
-from typing import Optional
+
+import typing
 
 from mitmproxy.contentviews import base
 
+PARSE_ERROR = object()
+
 
-def pretty_json(s: bytes) -> Optional[bytes]:
+def parse_json(s: bytes) -> typing.Any:
     try:
-        p = json.loads(s.decode('utf-8'))
+        return json.loads(s.decode('utf-8'))
     except ValueError:
-        return None
-    return p
+        return PARSE_ERROR
+
+
+def format_json(data: typing.Any) -> typing.Iterator[base.TViewLine]:
+    encoder = json.JSONEncoder(indent=4, sort_keys=True, ensure_ascii=False)
+    current_line: base.TViewLine = []
+    for chunk in encoder.iterencode(data):
+        if "\n" in chunk:
+            rest_of_last_line, chunk = chunk.split("\n", maxsplit=1)
+            # rest_of_last_line is a delimiter such as , or [
+            current_line.append(('text', rest_of_last_line))
+            yield current_line
+            current_line = []
+        if re.match(r'\s*"', chunk):
+            current_line.append(('json_string', chunk))
+        elif re.match(r'\s*\d', chunk):
+            current_line.append(('json_number', chunk))
+        elif re.match(r'\s*(true|null|false)', chunk):
+            current_line.append(('json_boolean', chunk))
+        else:
+            current_line.append(('text', chunk))
+    yield current_line
 
 
 class ViewJSON(base.View):
@@ -21,30 +44,7 @@ class ViewJSON(base.View):
         "application/vnd.api+json"
     ]
 
-    @staticmethod
-    def _format(pj):
-        li = []
-        for chunk in json.JSONEncoder(indent=4, sort_keys=True, ensure_ascii=False).iterencode(pj):
-            k = re.split("\\n", chunk)
-            if(len(k) > 1):
-                if(len(k[0]) > 0):
-                    li.append(('text', k[0]))
-                yield li
-                li = []
-                chunk = k[1]
-            else:
-                chunk = k[0]
-            if(re.match('^\s*\".*\"$', chunk)):
-                li.append(('json_string', chunk))
-            elif(re.match('\s*[0-9]+[.]{0,1}[0-9]*', chunk)):
-                li.append(('json_number', chunk))
-            elif(re.match('\s*true|null|false', chunk)):
-                li.append(('json_boolean', chunk))
-            else:
-                li.append(('text', chunk))
-        yield li
-
     def __call__(self, data, **metadata):
-        pj = pretty_json(data)
-        if pj is not None:
-            return "JSON", self._format(pj)
+        data = parse_json(data)
+        if data is not PARSE_ERROR:
+            return "JSON", format_json(data)
diff --git a/test/mitmproxy/contentviews/test_json.py b/test/mitmproxy/contentviews/test_json.py
index 3f522b7f..2b5bf86a 100644
--- a/test/mitmproxy/contentviews/test_json.py
+++ b/test/mitmproxy/contentviews/test_json.py
@@ -1,18 +1,43 @@
+from hypothesis import given
+from hypothesis.strategies import binary
+
 from mitmproxy.contentviews import json
 from . import full_eval
 
 
-def test_pretty_json():
-    assert json.pretty_json(b'{"foo": 1}')
-    assert not json.pretty_json(b"moo")
-    assert json.pretty_json(b'{"foo" : "\xe4\xb8\x96\xe7\x95\x8c"}')  # utf8 with chinese characters
-    assert not json.pretty_json(b'{"foo" : "\xFF"}')
+def test_parse_json():
+    assert json.parse_json(b'{"foo": 1}')
+    assert json.parse_json(b'null') is None
+    assert json.parse_json(b"moo") is json.PARSE_ERROR
+    assert json.parse_json(b'{"foo" : "\xe4\xb8\x96\xe7\x95\x8c"}')  # utf8 with chinese characters
+    assert json.parse_json(b'{"foo" : "\xFF"}') is json.PARSE_ERROR
+
+
+def test_format_json():
+    assert list(json.format_json({
+        "data": [
+            "str",
+            42,
+            True,
+            False,
+            None,
+            {},
+            []
+        ]
+    }))
 
 
 def test_view_json():
     v = full_eval(json.ViewJSON())
+    assert v(b"null")
     assert v(b"{}")
     assert not v(b"{")
     assert v(b"[1, 2, 3, 4, 5]")
     assert v(b'{"foo" : 3}')
     assert v(b'{"foo": true, "nullvalue": null}')
+
+
+@given(binary())
+def test_view_json_doesnt_crash(data):
+    v = full_eval(json.ViewJSON())
+    v(data)