mypy all of the codebase bar tnetstring

In some places, this involved removing type declarations where our types were
terminally confused. The grideditor specifically needs a cleanup and
restructure.

14 files changed, 77 insertions, 55 deletions

diff --git a/mitmproxy/addons/view.py b/mitmproxy/addons/view.py
index 1b8a30e4..7e9d66a1 100644
--- a/mitmproxy/addons/view.py
+++ b/mitmproxy/addons/view.py
@@ -18,6 +18,7 @@ import sortedcontainers
 import mitmproxy.flow
 from mitmproxy import flowfilter
 from mitmproxy import exceptions
+from mitmproxy import http  # noqa
 
 # The underlying sorted list implementation expects the sort key to be stable
 # for the lifetime of the object. However, if we sort by size, for instance,
@@ -34,7 +35,7 @@ class _OrderKey:
     def __init__(self, view):
         self.view = view
 
-    def generate(self, f: mitmproxy.flow.Flow) -> typing.Any:  # pragma: no cover
+    def generate(self, f: http.HTTPFlow) -> typing.Any:  # pragma: no cover
         pass
 
     def refresh(self, f):
@@ -64,22 +65,22 @@ class _OrderKey:
 
 
 class OrderRequestStart(_OrderKey):
-    def generate(self, f: mitmproxy.flow.Flow) -> datetime.datetime:
+    def generate(self, f: http.HTTPFlow) -> datetime.datetime:
         return f.request.timestamp_start or 0
 
 
 class OrderRequestMethod(_OrderKey):
-    def generate(self, f: mitmproxy.flow.Flow) -> str:
+    def generate(self, f: http.HTTPFlow) -> str:
         return f.request.method
 
 
 class OrderRequestURL(_OrderKey):
-    def generate(self, f: mitmproxy.flow.Flow) -> str:
+    def generate(self, f: http.HTTPFlow) -> str:
         return f.request.url
 
 
 class OrderKeySize(_OrderKey):
-    def generate(self, f: mitmproxy.flow.Flow) -> int:
+    def generate(self, f: http.HTTPFlow) -> int:
         s = 0
         if f.request.raw_content:
             s += len(f.request.raw_content)
@@ -118,7 +119,9 @@ class View(collections.Sequence):
         self.order_reversed = False
         self.focus_follow = False
 
-        self._view = sortedcontainers.SortedListWithKey(key = self.order_key)
+        self._view = sortedcontainers.SortedListWithKey(
+            key = self.order_key
+        )
 
         # The sig_view* signals broadcast events that affect the view. That is,
         # an update to a flow in the store but not in the view does not trigger
@@ -165,7 +168,7 @@ class View(collections.Sequence):
     def __len__(self):
         return len(self._view)
 
-    def __getitem__(self, offset) -> mitmproxy.flow.Flow:
+    def __getitem__(self, offset) -> typing.Any:
         return self._view[self._rev(offset)]
 
     # Reflect some methods to the efficient underlying implementation
@@ -177,7 +180,7 @@ class View(collections.Sequence):
     def index(self, f: mitmproxy.flow.Flow, start: int = 0, stop: typing.Optional[int] = None) -> int:
         return self._rev(self._view.index(f, start, stop))
 
-    def __contains__(self, f: mitmproxy.flow.Flow) -> bool:
+    def __contains__(self, f: typing.Any) -> bool:
         return self._view.__contains__(f)
 
     def _order_key_name(self):
@@ -402,7 +405,7 @@ class Focus:
 class Settings(collections.Mapping):
     def __init__(self, view: View) -> None:
         self.view = view
-        self._values = {}  # type: typing.MutableMapping[str, mitmproxy.flow.Flow]
+        self._values = {}  # type: typing.MutableMapping[str, typing.Dict]
         view.sig_store_remove.connect(self._sig_store_remove)
         view.sig_store_refresh.connect(self._sig_store_refresh)
 
diff --git a/mitmproxy/flowfilter.py b/mitmproxy/flowfilter.py
index 2c7fc52f..83c98bad 100644
--- a/mitmproxy/flowfilter.py
+++ b/mitmproxy/flowfilter.py
@@ -44,7 +44,7 @@ from mitmproxy import flow
 from mitmproxy.utils import strutils
 
 import pyparsing as pp
-from typing import Callable
+from typing import Callable, Sequence, Type  # noqa
 
 
 def only(*types):
@@ -69,6 +69,8 @@ class _Token:
 
 
 class _Action(_Token):
+    code = None  # type: str
+    help = None  # type: str
 
     @classmethod
     def make(klass, s, loc, toks):
@@ -162,15 +164,14 @@ def _check_content_type(rex, message):
 class FAsset(_Action):
     code = "a"
     help = "Match asset in response: CSS, Javascript, Flash, images."
-    ASSET_TYPES = [
+    ASSET_TYPES = [re.compile(x) for x in [
         b"text/javascript",
         b"application/x-javascript",
         b"application/javascript",
         b"text/css",
         b"image/.*",
         b"application/x-shockwave-flash"
-    ]
-    ASSET_TYPES = [re.compile(x) for x in ASSET_TYPES]
+    ]]
 
     @only(http.HTTPFlow)
     def __call__(self, f):
@@ -436,7 +437,7 @@ filter_unary = [
     FResp,
     FTCP,
     FWebSocket,
-]
+]  # type: Sequence[Type[_Action]]
 filter_rex = [
     FBod,
     FBodRequest,
@@ -452,7 +453,7 @@ filter_rex = [
     FMethod,
     FSrc,
     FUrl,
-]
+]  # type: Sequence[Type[_Rex]]
 filter_int = [
     FCode
 ]
@@ -538,17 +539,17 @@ def match(flt, flow):
 
 
 help = []
-for i in filter_unary:
+for a in filter_unary:
     help.append(
-        ("~%s" % i.code, i.help)
+        ("~%s" % a.code, a.help)
     )
-for i in filter_rex:
+for b in filter_rex:
     help.append(
-        ("~%s regex" % i.code, i.help)
+        ("~%s regex" % b.code, b.help)
     )
-for i in filter_int:
+for c in filter_int:
     help.append(
-        ("~%s int" % i.code, i.help)
+        ("~%s int" % c.code, c.help)
     )
 help.sort()
 help.extend(
diff --git a/mitmproxy/io.py b/mitmproxy/io.py
index 780955a4..0f6c3f5c 100644
--- a/mitmproxy/io.py
+++ b/mitmproxy/io.py
@@ -1,5 +1,5 @@
 import os
-from typing import Iterable
+from typing import Type, Iterable, Dict, Union, Any, cast  # noqa
 
 from mitmproxy import exceptions
 from mitmproxy import flow
@@ -15,7 +15,7 @@ FLOW_TYPES = dict(
     http=http.HTTPFlow,
     websocket=websocket.WebSocketFlow,
     tcp=tcp.TCPFlow,
-)
+)  # type: Dict[str, Type[flow.Flow]]
 
 
 class FlowWriter:
@@ -37,14 +37,18 @@ class FlowReader:
         """
         try:
             while True:
-                data = tnetstring.load(self.fo)
+                # FIXME: This cast hides a lack of dynamic type checking
+                loaded = cast(
+                    Dict[Union[bytes, str], Any],
+                    tnetstring.load(self.fo),
+                )
                 try:
-                    data = io_compat.migrate_flow(data)
+                    mdata = io_compat.migrate_flow(loaded)
                 except ValueError as e:
                     raise exceptions.FlowReadException(str(e))
-                if data["type"] not in FLOW_TYPES:
-                    raise exceptions.FlowReadException("Unknown flow type: {}".format(data["type"]))
-                yield FLOW_TYPES[data["type"]].from_state(data)
+                if mdata["type"] not in FLOW_TYPES:
+                    raise exceptions.FlowReadException("Unknown flow type: {}".format(mdata["type"]))
+                yield FLOW_TYPES[mdata["type"]].from_state(mdata)
         except ValueError as e:
             if str(e) == "not a tnetstring: empty file":
                 return  # Error is due to EOF
diff --git a/mitmproxy/io_compat.py b/mitmproxy/io_compat.py
index 7d839ffd..9d95f602 100644
--- a/mitmproxy/io_compat.py
+++ b/mitmproxy/io_compat.py
@@ -2,7 +2,7 @@
 This module handles the import of mitmproxy flows generated by old versions.
 """
 import uuid
-from typing import Any, Dict
+from typing import Any, Dict, Mapping, Union  # noqa
 
 from mitmproxy import version
 from mitmproxy.utils import strutils
@@ -113,8 +113,8 @@ def convert_300_4(data):
     return data
 
 
-client_connections = {}
-server_connections = {}
+client_connections = {}  # type: Mapping[str, str]
+server_connections = {}  # type: Mapping[str, str]
 
 
 def convert_4_5(data):
@@ -187,7 +187,7 @@ converters = {
 }
 
 
-def migrate_flow(flow_data: Dict[str, Any]) -> Dict[str, Any]:
+def migrate_flow(flow_data: Dict[Union[bytes, str], Any]) -> Dict[Union[bytes, str], Any]:
     while True:
         flow_version = flow_data.get(b"version", flow_data.get("version"))
 
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index 04052ec8..45377b2c 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -5,6 +5,7 @@ from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
 from mitmproxy.addons import view
 from mitmproxy import export
+import mitmproxy.tools.console.master # noqa
 
 
 def _mkhelp():
@@ -305,7 +306,9 @@ class FlowListWalker(urwid.ListWalker):
 
 class FlowListBox(urwid.ListBox):
 
-    def __init__(self, master: "mitmproxy.tools.console.master.ConsoleMaster"):
+    def __init__(
+        self, master: "mitmproxy.tools.console.master.ConsoleMaster"
+   ) -> None:
         self.master = master  # type: "mitmproxy.tools.console.master.ConsoleMaster"
         super().__init__(FlowListWalker(master))
 
diff --git a/mitmproxy/tools/console/flowview.py b/mitmproxy/tools/console/flowview.py
index ba41c947..33c8f2ac 100644
--- a/mitmproxy/tools/console/flowview.py
+++ b/mitmproxy/tools/console/flowview.py
@@ -19,6 +19,7 @@ from mitmproxy.tools.console import overlay
 from mitmproxy.tools.console import searchable
 from mitmproxy.tools.console import signals
 from mitmproxy.tools.console import tabs
+import mitmproxy.tools.console.master # noqa
 
 
 class SearchError(Exception):
@@ -103,7 +104,11 @@ footer = [
 
 class FlowViewHeader(urwid.WidgetWrap):
 
-    def __init__(self, master: "mitmproxy.console.master.ConsoleMaster", f: http.HTTPFlow):
+    def __init__(
+        self,
+        master: "mitmproxy.tools.console.master.ConsoleMaster",
+        f: http.HTTPFlow
+    ) -> None:
         self.master = master
         self.flow = f
         self._w = common.format_flow(
@@ -651,8 +656,8 @@ class FlowView(tabs.Tabs):
                     )
         elif key == "z":
             self.flow.backup()
-            e = conn.headers.get("content-encoding", "identity")
-            if e != "identity":
+            enc = conn.headers.get("content-encoding", "identity")
+            if enc != "identity":
                 try:
                     conn.decode()
                 except ValueError:
diff --git a/mitmproxy/tools/console/grideditor/base.py b/mitmproxy/tools/console/grideditor/base.py
index d2ba47c3..151479a4 100644
--- a/mitmproxy/tools/console/grideditor/base.py
+++ b/mitmproxy/tools/console/grideditor/base.py
@@ -7,10 +7,12 @@ from typing import Iterable
 from typing import Optional
 from typing import Sequence
 from typing import Tuple
+from typing import Set # noqa
 
 import urwid
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
+import mitmproxy.tools.console.master # noqa
 
 FOOTER = [
     ('heading_key', "enter"), ":edit ",
@@ -34,7 +36,7 @@ class Cell(urwid.WidgetWrap):
 
 
 class Column(metaclass=abc.ABCMeta):
-    subeditor = None
+    subeditor = None  # type: urwid.Edit
 
     def __init__(self, heading):
         self.heading = heading
@@ -62,13 +64,13 @@ class GridRow(urwid.WidgetWrap):
             editing: bool,
             editor: "GridEditor",
             values: Tuple[Iterable[bytes], Container[int]]
-    ):
+    ) -> None:
         self.focused = focused
         self.editor = editor
         self.edit_col = None  # type: Optional[Cell]
 
         errors = values[1]
-        self.fields = []
+        self.fields = []  # type: Sequence[Any]
         for i, v in enumerate(values[0]):
             if focused == i and editing:
                 self.edit_col = self.editor.columns[i].Edit(v)
@@ -116,8 +118,8 @@ class GridWalker(urwid.ListWalker):
             self,
             lst: Iterable[list],
             editor: "GridEditor"
-    ):
-        self.lst = [(i, set()) for i in lst]
+    ) -> None:
+        self.lst = [(i, set()) for i in lst]  # type: Sequence[Tuple[Any, Set]]
         self.editor = editor
         self.focus = 0
         self.focus_col = 0
@@ -256,12 +258,12 @@ class GridEditor(urwid.WidgetWrap):
 
     def __init__(
             self,
-            master: "mitmproxy.console.master.ConsoleMaster",
+            master: "mitmproxy.tools.console.master.ConsoleMaster",
             value: Any,
             callback: Callable[..., None],
             *cb_args,
             **cb_kwargs
-    ):
+    ) -> None:
         value = self.data_in(copy.deepcopy(value))
         self.master = master
         self.value = value
@@ -380,7 +382,7 @@ class GridEditor(urwid.WidgetWrap):
         """
             Return None, or a string error message.
         """
-        return False
+        return None
 
     def handle_key(self, key):
         return False
diff --git a/mitmproxy/tools/console/grideditor/col_bytes.py b/mitmproxy/tools/console/grideditor/col_bytes.py
index f580e947..e4a53453 100644
--- a/mitmproxy/tools/console/grideditor/col_bytes.py
+++ b/mitmproxy/tools/console/grideditor/col_bytes.py
@@ -9,7 +9,7 @@ from mitmproxy.utils import strutils
 
 def read_file(filename: str, callback: Callable[..., None], escaped: bool) -> Optional[str]:
     if not filename:
-        return
+        return None
 
     filename = os.path.expanduser(filename)
     try:
@@ -26,6 +26,7 @@ def read_file(filename: str, callback: Callable[..., None], escaped: bool) -> Op
     # TODO: Refactor the status_prompt_path signal so that we
     # can raise exceptions here and return the content instead.
     callback(d)
+    return None
 
 
 class Column(base.Column):
@@ -68,7 +69,7 @@ class Column(base.Column):
 
 
 class Display(base.Cell):
-    def __init__(self, data: bytes):
+    def __init__(self, data: bytes) -> None:
         self.data = data
         escaped = strutils.bytes_to_escaped_str(data)
         w = urwid.Text(escaped, wrap="any")
@@ -79,7 +80,7 @@ class Display(base.Cell):
 
 
 class Edit(base.Cell):
-    def __init__(self, data: bytes):
+    def __init__(self, data: bytes) -> None:
         data = strutils.bytes_to_escaped_str(data)
         w = urwid.Edit(edit_text=data, wrap="any", multiline=True)
         w = urwid.AttrWrap(w, "editfield")
diff --git a/mitmproxy/tools/console/grideditor/col_text.py b/mitmproxy/tools/console/grideditor/col_text.py
index 430ad037..f0ac06f8 100644
--- a/mitmproxy/tools/console/grideditor/col_text.py
+++ b/mitmproxy/tools/console/grideditor/col_text.py
@@ -26,12 +26,11 @@ class Column(col_bytes.Column):
 
 # This is the same for both edit and display.
 class EncodingMixin:
-    def __init__(self, data: str, encoding_args) -> "TDisplay":
+    def __init__(self, data, encoding_args):
         self.encoding_args = encoding_args
-        data = data.encode(*self.encoding_args)
-        super().__init__(data)
+        super().__init__(data.encode(*self.encoding_args))
 
-    def get_data(self) -> str:
+    def get_data(self):
         data = super().get_data()
         try:
             return data.decode(*self.encoding_args)
diff --git a/mitmproxy/tools/console/grideditor/editors.py b/mitmproxy/tools/console/grideditor/editors.py
index 313495e4..e069fe2f 100644
--- a/mitmproxy/tools/console/grideditor/editors.py
+++ b/mitmproxy/tools/console/grideditor/editors.py
@@ -248,7 +248,7 @@ class SetCookieEditor(base.GridEditor):
 
 
 class OptionsEditor(base.GridEditor):
-    title = None
+    title = None  # type: str
     columns = [
         col_text.Column("")
     ]
diff --git a/mitmproxy/tools/console/master.py b/mitmproxy/tools/console/master.py
index c0d8e05c..c1d584ac 100644
--- a/mitmproxy/tools/console/master.py
+++ b/mitmproxy/tools/console/master.py
@@ -70,7 +70,6 @@ class UnsupportedLog:
 
 
 class ConsoleMaster(master.Master):
-    palette = []
 
     def __init__(self, options, server):
         super().__init__(options, server)
diff --git a/mitmproxy/tools/console/palettes.py b/mitmproxy/tools/console/palettes.py
index 7b15f98f..e349f702 100644
--- a/mitmproxy/tools/console/palettes.py
+++ b/mitmproxy/tools/console/palettes.py
@@ -1,3 +1,4 @@
+import typing
 # Low-color themes should ONLY use the standard foreground and background
 # colours listed here:
 #
@@ -32,7 +33,7 @@ class Palette:
         # Grid Editor
         'focusfield', 'focusfield_error', 'field_error', 'editfield',
     ]
-    high = None
+    high = None  # type: typing.Mapping[str, typing.Sequence[str]]
 
     def palette(self, transparent):
         l = []
diff --git a/mitmproxy/tools/console/statusbar.py b/mitmproxy/tools/console/statusbar.py
index c7132864..d3a3e1f2 100644
--- a/mitmproxy/tools/console/statusbar.py
+++ b/mitmproxy/tools/console/statusbar.py
@@ -5,6 +5,7 @@ import urwid
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import pathedit
 from mitmproxy.tools.console import signals
+import mitmproxy.tools.console.master # noqa
 
 
 class PromptPath:
@@ -135,7 +136,9 @@ class ActionBar(urwid.WidgetWrap):
 
 class StatusBar(urwid.WidgetWrap):
 
-    def __init__(self, master: "mitmproxy.console.master.ConsoleMaster", helptext):
+    def __init__(
+        self, master: "mitmproxy.tools.console.master.ConsoleMaster", helptext
+    ) -> None:
         self.master = master
         self.helptext = helptext
         self.ib = urwid.WidgetWrap(urwid.Text(""))
diff --git a/mitmproxy/tools/web/app.py b/mitmproxy/tools/web/app.py
index 002513b9..23d620e0 100644
--- a/mitmproxy/tools/web/app.py
+++ b/mitmproxy/tools/web/app.py
@@ -17,6 +17,7 @@ from mitmproxy import http
 from mitmproxy import io
 from mitmproxy import log
 from mitmproxy import version
+import mitmproxy.tools.web.master # noqa
 
 
 def flow_to_json(flow: mitmproxy.flow.Flow) -> dict: