Merge pull request #2190 from cortesi/mypy

mypy for all of ./mitmproxy

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)
 

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(

mitmproxy/io/__init__.py

@@ -0,0 +1,7 @@
+
+from .io import FlowWriter, FlowReader, FilteredFlowWriter, read_flows_from_paths
+
+
+__all__ = [
+    "FlowWriter", "FlowReader", "FilteredFlowWriter", "read_flows_from_paths"
+]

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"))
 

mitmproxy/io/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
@@ -7,15 +7,15 @@ from mitmproxy import flowfilter
 from mitmproxy import http
 from mitmproxy import tcp
 from mitmproxy import websocket
-from mitmproxy.contrib import tnetstring
-from mitmproxy import io_compat
 
+from mitmproxy.io import compat
+from mitmproxy.io import tnetstring
 
 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 = 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

mitmproxy/io/tnetstring.py

@@ -41,9 +41,9 @@ all other strings are returned as plain bytes.
 """
 
 import collections
-from typing import io, Union, Tuple
+import typing
 
-TSerializable = Union[None, bool, int, float, bytes, list, tuple, dict]
+TSerializable = typing.Union[None, str, bool, int, float, bytes, list, tuple, dict]
 
 
 def dumps(value: TSerializable) -> bytes:
@@ -53,12 +53,12 @@ def dumps(value: TSerializable) -> bytes:
     #  This uses a deque to collect output fragments in reverse order,
     #  then joins them together at the end.  It's measurably faster
     #  than creating all the intermediate strings.
-    q = collections.deque()
+    q = collections.deque()  # type: collections.deque
     _rdumpq(q, 0, value)
     return b''.join(q)
 
 
-def dump(value: TSerializable, file_handle: io.BinaryIO) -> None:
+def dump(value: TSerializable, file_handle: typing.BinaryIO) -> None:
     """
     This function dumps a python object as a tnetstring and
     writes it to the given file.
@@ -156,7 +156,7 @@ def loads(string: bytes) -> TSerializable:
     return pop(string)[0]
 
 
-def load(file_handle: io.BinaryIO) -> TSerializable:
+def load(file_handle: typing.BinaryIO) -> TSerializable:
     """load(file) -> object
 
     This function reads a tnetstring from a file and parses it into a
@@ -213,19 +213,19 @@ def parse(data_type: int, data: bytes) -> TSerializable:
         l = []
         while data:
             item, data = pop(data)
-            l.append(item)
+            l.append(item)  # type: ignore
         return l
     if data_type == ord(b'}'):
         d = {}
         while data:
             key, data = pop(data)
             val, data = pop(data)
-            d[key] = val
+            d[key] = val  # type: ignore
         return d
     raise ValueError("unknown type tag: {}".format(data_type))
 
 
-def pop(data: bytes) -> Tuple[TSerializable, bytes]:
+def pop(data: bytes) -> typing.Tuple[TSerializable, bytes]:
     """
     This function parses a tnetstring into a python object.
     It returns a tuple giving the parsed object and a string
@@ -233,8 +233,8 @@ def pop(data: bytes) -> Tuple[TSerializable, bytes]:
     """
     #  Parse out data length, type and remaining string.
     try:
-        length, data = data.split(b':', 1)
-        length = int(length)
+        blength, data = data.split(b':', 1)
+        length = int(blength)
     except ValueError:
         raise ValueError("not a tnetstring: missing or invalid length prefix: {}".format(data))
     try:

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))
 

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:

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

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")

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)

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("")
     ]

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)

mitmproxy/tools/console/palettes.py

@@ -1,3 +1,4 @@
+import typing  # noqa
 # 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 = []

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(""))

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:

setup.cfg

@@ -37,7 +37,7 @@ exclude =
     mitmproxy/controller.py
     mitmproxy/export.py
     mitmproxy/flow.py
-    mitmproxy/io_compat.py
+    mitmproxy/io/compat.py
     mitmproxy/master.py
     pathod/pathoc.py
     pathod/pathod.py
@@ -57,8 +57,9 @@ exclude =
     mitmproxy/exceptions.py
     mitmproxy/export.py
     mitmproxy/flow.py
-    mitmproxy/io.py
-    mitmproxy/io_compat.py
+    mitmproxy/io/io.py
+    mitmproxy/io/compat.py
+    mitmproxy/io/tnetstring.py
     mitmproxy/log.py
     mitmproxy/master.py
     mitmproxy/net/check.py

test/mitmproxy/io/test_tnetstring.py

@@ -4,7 +4,7 @@ import math
 import io
 import struct
 
-from mitmproxy.contrib import tnetstring
+from mitmproxy.io import tnetstring
 
 MAXINT = 2 ** (struct.Struct('i').size * 8 - 1) - 1
 

test/mitmproxy/test_flow.py

@@ -6,7 +6,7 @@ import mitmproxy.io
 from mitmproxy import flowfilter
 from mitmproxy import options
 from mitmproxy.proxy import config
-from mitmproxy.contrib import tnetstring
+from mitmproxy.io import tnetstring
 from mitmproxy.exceptions import FlowReadException
 from mitmproxy import flow
 from mitmproxy import http

test/mitmproxy/test_websocket.py

@@ -1,7 +1,7 @@
 import io
 import pytest
 
-from mitmproxy.contrib import tnetstring
+from mitmproxy.io import tnetstring
 from mitmproxy import flowfilter
 from mitmproxy.test import tflow
 

tox.ini

@@ -27,17 +27,7 @@ commands =
   flake8 --jobs 8 mitmproxy pathod examples test release
   python3 test/filename_matching.py
   rstcheck README.rst
-  mypy --ignore-missing-imports --follow-imports=skip \
-    mitmproxy/addons/ \
-    mitmproxy/addonmanager.py \
-    mitmproxy/optmanager.py \
-    mitmproxy/proxy/protocol/ \
-    mitmproxy/log.py \
-    mitmproxy/tools/dump.py \
-    mitmproxy/tools/web/ \
-    mitmproxy/contentviews/
-  mypy --ignore-missing-imports \
-    mitmproxy/master.py
+  mypy --ignore-missing-imports ./mitmproxy
 
 [testenv:individual_coverage]
 deps =