update tests, increase coverage, add type info

8 files changed, 48 insertions, 45 deletions

diff --git a/mitmproxy/contentviews/__init__.py b/mitmproxy/contentviews/__init__.py
index 4c3cb99b..b83e7aa6 100644
--- a/mitmproxy/contentviews/__init__.py
+++ b/mitmproxy/contentviews/__init__.py
@@ -14,31 +14,37 @@ passed as the ``headers`` keyword argument. For HTTP requests, the query
 parameters are passed as the ``query`` keyword argument.
 """
 import traceback
+from typing import Dict, Optional  # noqa
+from typing import List  # noqa
+from typing import Tuple  # noqa
 
 from mitmproxy import exceptions
 from mitmproxy.net import http
 from mitmproxy.utils import strutils
+from . import (
+    auto, raw, hex, json, xml, wbxml, html, javascript, css,
+    urlencoded, multipart, image, query, protobuf
+)
+from .base import View, VIEW_CUTOFF, KEY_MAX, format_text, format_dict
 
-from .base import VIEW_CUTOFF, KEY_MAX
-
-views = []
-content_types_map = {}
-view_prompts = []
+views = []  # type: List[View]
+content_types_map = {}  # type: Dict[str, List[View]]
+view_prompts = []  # type: List[Tuple[str, str]]
 
 
-def get(name):
+def get(name: str) -> Optional[View]:
     for i in views:
         if i.name.lower() == name.lower():
             return i
 
 
-def get_by_shortcut(c):
+def get_by_shortcut(c: str) -> Optional[View]:
     for i in views:
         if i.prompt[1] == c:
             return i
 
 
-def add(view):
+def add(view: View) -> None:
     # TODO: auto-select a different name (append an integer?)
     for i in views:
         if i.name == view.name:
@@ -58,7 +64,7 @@ def add(view):
     view_prompts.append(view.prompt)
 
 
-def remove(view):
+def remove(view: View) -> None:
     for ct in view.content_types:
         l = content_types_map.setdefault(ct, [])
         l.remove(view)
@@ -123,7 +129,7 @@ def get_message_content_view(viewname, message):
     return description, lines, error
 
 
-def get_content_view(viewmode, data, **metadata):
+def get_content_view(viewmode: View, data: bytes, **metadata):
     """
         Args:
             viewmode: the view to use.
@@ -153,11 +159,6 @@ def get_content_view(viewmode, data, **metadata):
     return desc, safe_to_print(content), error
 
 
-from . import (
-    auto, raw, hex, json, xml, wbxml, html, javascript, css,
-    urlencoded, multipart, image, query, protobuf
-)
-
 add(auto.ViewAuto())
 add(raw.ViewRaw())
 add(hex.ViewHex())
@@ -174,4 +175,10 @@ add(image.ViewImage())
 add(query.ViewQuery())
 
 if protobuf.ViewProtobuf.is_available():
-    add(protobuf.ViewProtobuf())
\ No newline at end of file
+    add(protobuf.ViewProtobuf())
+
+__all__ = [
+    "View", "VIEW_CUTOFF", "KEY_MAX", "format_text", "format_dict",
+    "get", "get_by_shortcut", "add", "remove",
+    "get_content_view", "get_message_content_view",
+]
diff --git a/mitmproxy/contentviews/auto.py b/mitmproxy/contentviews/auto.py
index 2b08f165..3a162f53 100644
--- a/mitmproxy/contentviews/auto.py
+++ b/mitmproxy/contentviews/auto.py
@@ -1,12 +1,12 @@
+from mitmproxy import contentviews
 from mitmproxy.net import http
 from mitmproxy.utils import strutils
 from . import base
-from mitmproxy.contentviews import get, content_types_map
+
 
 class ViewAuto(base.View):
     name = "Auto"
     prompt = ("auto", "a")
-    content_types = []
 
     def __call__(self, data, **metadata):
         headers = metadata.get("headers", {})
@@ -14,14 +14,14 @@ class ViewAuto(base.View):
         if data and ctype:
             ct = http.parse_content_type(ctype) if ctype else None
             ct = "%s/%s" % (ct[0], ct[1])
-            if ct in content_types_map:
-                return content_types_map[ct][0](data, **metadata)
+            if ct in contentviews.content_types_map:
+                return contentviews.content_types_map[ct][0](data, **metadata)
             elif strutils.is_xml(data):
-                return get("XML")(data, **metadata)
+                return contentviews.get("XML")(data, **metadata)
         if metadata.get("query"):
-            return get("Query")(data, **metadata)
+            return contentviews.get("Query")(data, **metadata)
         if data and strutils.is_mostly_bin(data):
-            return get("Hex")(data)
+            return contentviews.get("Hex")(data)
         if not data:
             return "No content", []
-        return get("Raw")(data)
+        return contentviews.get("Raw")(data)
diff --git a/mitmproxy/contentviews/base.py b/mitmproxy/contentviews/base.py
index b1a51ffe..0de4f786 100644
--- a/mitmproxy/contentviews/base.py
+++ b/mitmproxy/contentviews/base.py
@@ -1,9 +1,8 @@
 # Default view cutoff *in lines*
 
-from typing import Iterable
+from typing import Iterable, AnyStr, List
 from typing import Mapping
 from typing import Tuple
-from typing import Union
 
 VIEW_CUTOFF = 512
 
@@ -11,9 +10,9 @@ KEY_MAX = 30
 
 
 class View:
-    name = None
-    prompt = ()
-    content_types = []
+    name = None  # type: str
+    prompt = None  # type: Tuple[str,str]
+    content_types = []  # type: List[str]
 
     def __call__(self, data: bytes, **metadata):
         """
@@ -35,12 +34,12 @@ class View:
             The content generator must not yield tuples of tuples,
             because urwid cannot process that. You have to yield a *list* of tuples per line.
         """
-        raise NotImplementedError()
+        raise NotImplementedError()  # pragma: no cover
 
 
 def format_dict(
-        d: Mapping[Union[str, bytes], Union[str, bytes]]
-) -> Iterable[Tuple[Union[str, bytes], Union[str, bytes]]]:
+        d: Mapping[AnyStr, AnyStr]
+) -> Iterable[List[Tuple[str, AnyStr]]]:
     """
     Helper function that transforms the given dictionary into a list of
         ("key",   key  )
@@ -58,7 +57,7 @@ def format_dict(
         ]
 
 
-def format_text(text):
+def format_text(text: AnyStr) -> Iterable[List[Tuple[str, AnyStr]]]:
     """
     Helper function that transforms bytes into the view output format.
     """
diff --git a/mitmproxy/contentviews/hex.py b/mitmproxy/contentviews/hex.py
index 116ed600..6251a8f3 100644
--- a/mitmproxy/contentviews/hex.py
+++ b/mitmproxy/contentviews/hex.py
@@ -5,7 +5,6 @@ from . import base
 class ViewHex(base.View):
     name = "Hex"
     prompt = ("hex", "e")
-    content_types = []
 
     @staticmethod
     def _format(data):
diff --git a/mitmproxy/contentviews/json.py b/mitmproxy/contentviews/json.py
index 7c128d02..de7f1093 100644
--- a/mitmproxy/contentviews/json.py
+++ b/mitmproxy/contentviews/json.py
@@ -1,7 +1,7 @@
 import json
 from typing import Optional
 
-from mitmproxy.contentviews.base import format_text, View
+from mitmproxy.contentviews import base
 
 
 def pretty_json(s: bytes) -> Optional[bytes]:
@@ -10,15 +10,10 @@ def pretty_json(s: bytes) -> Optional[bytes]:
     except ValueError:
         return None
     pretty = json.dumps(p, sort_keys=True, indent=4, ensure_ascii=False)
-    if isinstance(pretty, str):
-        # json.dumps _may_ decide to return unicode, if the JSON object is not ascii.
-        # From limited testing this is always valid utf8 (otherwise json.loads will fail earlier),
-        # so we can just re-encode it here.
-        return pretty.encode("utf8", "strict")
-    return pretty
+    return pretty.encode("utf8", "strict")
 
 
-class ViewJSON(View):
+class ViewJSON(base.View):
     name = "JSON"
     prompt = ("json", "s")
     content_types = [
@@ -29,4 +24,4 @@ class ViewJSON(View):
     def __call__(self, data, **metadata):
         pj = pretty_json(data)
         if pj:
-            return "JSON", format_text(pj)
+            return "JSON", base.format_text(pj)
diff --git a/mitmproxy/contentviews/multipart.py b/mitmproxy/contentviews/multipart.py
index 640896ab..0b0e51e2 100644
--- a/mitmproxy/contentviews/multipart.py
+++ b/mitmproxy/contentviews/multipart.py
@@ -2,6 +2,7 @@ from mitmproxy.net import http
 from mitmproxy.types import multidict
 from . import base
 
+
 class ViewMultipart(base.View):
     name = "Multipart Form"
     prompt = ("multipart", "m")
diff --git a/mitmproxy/contentviews/query.py b/mitmproxy/contentviews/query.py
index c4ce0faf..93f47829 100644
--- a/mitmproxy/contentviews/query.py
+++ b/mitmproxy/contentviews/query.py
@@ -1,10 +1,11 @@
+from typing import List  # noqa
+
 from . import base
 
 
 class ViewQuery(base.View):
     name = "Query"
     prompt = ("query", "q")
-    content_types = []
 
     def __call__(self, data, **metadata):
         query = metadata.get("query")
diff --git a/mitmproxy/contentviews/raw.py b/mitmproxy/contentviews/raw.py
index c504a461..dcc53aa7 100644
--- a/mitmproxy/contentviews/raw.py
+++ b/mitmproxy/contentviews/raw.py
@@ -1,3 +1,5 @@
+from typing import List  # noqa
+
 from mitmproxy.utils import strutils
 from . import base
 
@@ -5,7 +7,6 @@ from . import base
 class ViewRaw(base.View):
     name = "Raw"
     prompt = ("raw", "r")
-    content_types = []
 
     def __call__(self, data, **metadata):
         return "Raw", base.format_text(strutils.bytes_to_escaped_str(data, True))