stateobject: use typing, enable tuples and more complex datatypes

mitmproxy/flow.py

@@ -87,7 +87,7 @@ def __init__(
         type=str,
         intercepted=bool,
         marked=bool,
-        metadata=dict,
+        metadata=typing.Dict[str, typing.Any],
     )
 
     def get_state(self):

mitmproxy/stateobject.py

@@ -1,18 +1,12 @@
-from typing import Any
-from typing import List
+import typing
+from typing import Any  # noqa
 from typing import MutableMapping  # noqa
 
 from mitmproxy.coretypes import serializable
-
-
-def _is_list(cls):
-    # The typing module is broken on Python 3.5.0, fixed on 3.5.1.
-    is_list_bugfix = getattr(cls, "__origin__", False) == getattr(List[Any], "__origin__", True)
-    return issubclass(cls, List) or is_list_bugfix
+from mitmproxy.utils import typecheck
 
 
 class StateObject(serializable.Serializable):
-
     """
     An object with serializable state.
 
@@ -34,22 +28,7 @@ def get_state(self):
         state = {}
         for attr, cls in self._stateobject_attributes.items():
             val = getattr(self, attr)
-            if val is None:
-                state[attr] = None
-            elif hasattr(val, "get_state"):
-                state[attr] = val.get_state()
-            elif _is_list(cls):
-                state[attr] = [x.get_state() for x in val]
-            elif isinstance(val, dict):
-                s = {}
-                for k, v in val.items():
-                    if hasattr(v, "get_state"):
-                        s[k] = v.get_state()
-                    else:
-                        s[k] = v
-                state[attr] = s
-            else:
-                state[attr] = val
+            state[attr] = get_state(cls, val)
         return state
 
     def set_state(self, state):
@@ -65,13 +44,51 @@ def set_state(self, state):
                 curr = getattr(self, attr)
                 if hasattr(curr, "set_state"):
                     curr.set_state(val)
-                elif hasattr(cls, "from_state"):
-                    obj = cls.from_state(val)
-                    setattr(self, attr, obj)
-                elif _is_list(cls):
-                    cls = cls.__parameters__[0] if cls.__parameters__ else cls.__args__[0]
-                    setattr(self, attr, [cls.from_state(x) for x in val])
-                else:  # primitive types such as int, str, ...
-                    setattr(self, attr, cls(val))
+                else:
+                    setattr(self, attr, make_object(cls, val))
         if state:
             raise RuntimeWarning("Unexpected State in __setstate__: {}".format(state))
+
+
+def _process(typeinfo: typecheck.Type, val: typing.Any, make: bool) -> typing.Any:
+    if val is None:
+        return None
+    elif make and hasattr(typeinfo, "from_state"):
+        return typeinfo.from_state(val)
+    elif not make and hasattr(val, "get_state"):
+        return val.get_state()
+
+    typename = str(typeinfo)
+
+    if typename.startswith("typing.List"):
+        T = typecheck.sequence_type(typeinfo)
+        return [_process(T, x, make) for x in val]
+    elif typename.startswith("typing.Tuple"):
+        Ts = typecheck.tuple_types(typeinfo)
+        if len(Ts) != len(val):
+            raise ValueError("Invalid data. Expected {}, got {}.".format(Ts, val))
+        return tuple(
+            _process(T, x, make) for T, x in zip(Ts, val)
+        )
+    elif typename.startswith("typing.Dict"):
+        k_cls, v_cls = typecheck.mapping_types(typeinfo)
+        return {
+            _process(k_cls, k, make): _process(v_cls, v, make)
+            for k, v in val.items()
+        }
+    elif typename.startswith("typing.Any"):
+        # FIXME: Remove this when we remove flow.metadata
+        assert isinstance(val, (int, str, bool, bytes))
+        return val
+    else:
+        return typeinfo(val)
+
+
+def make_object(typeinfo: typecheck.Type, val: typing.Any) -> typing.Any:
+    """Create an object based on the state given in val."""
+    return _process(typeinfo, val, True)
+
+
+def get_state(typeinfo: typecheck.Type, val: typing.Any) -> typing.Any:
+    """Get the state of the object given as val."""
+    return _process(typeinfo, val, False)

mitmproxy/utils/typecheck.py

@@ -1,7 +1,40 @@
 import typing
 
+Type = typing.Union[
+    typing.Any  # anything more elaborate really fails with mypy at the moment.
+]
 
-def check_option_type(name: str, value: typing.Any, typeinfo: typing.Any) -> None:
+
+def sequence_type(typeinfo: typing.Type[typing.List]) -> Type:
+    """Return the type of a sequence, e.g. typing.List"""
+    try:
+        return typeinfo.__args__[0]  # type: ignore
+    except AttributeError:  # Python 3.5.0
+        return typeinfo.__parameters__[0]  # type: ignore
+
+
+def tuple_types(typeinfo: typing.Type[typing.Tuple]) -> typing.Sequence[Type]:
+    """Return the types of a typing.Tuple"""
+    try:
+        return typeinfo.__args__  # type: ignore
+    except AttributeError:  # Python 3.5.x
+        return typeinfo.__tuple_params__  # type: ignore
+
+
+def union_types(typeinfo: typing.Type[typing.Tuple]) -> typing.Sequence[Type]:
+    """return the types of a typing.Union"""
+    try:
+        return typeinfo.__args__  # type: ignore
+    except AttributeError:  # Python 3.5.x
+        return typeinfo.__union_params__  # type: ignore
+
+
+def mapping_types(typeinfo: typing.Type[typing.Mapping]) -> typing.Tuple[Type, Type]:
+    """return the types of a mapping, e.g. typing.Dict"""
+    return typeinfo.__args__  # type: ignore
+
+
+def check_option_type(name: str, value: typing.Any, typeinfo: Type) -> None:
     """
     Check if the provided value is an instance of typeinfo and raises a
     TypeError otherwise. This function supports only those types required for
@@ -16,13 +49,7 @@ def check_option_type(name: str, value: typing.Any, typeinfo: typing.Any) -> Non
     typename = str(typeinfo)
 
     if typename.startswith("typing.Union"):
-        try:
-            types = typeinfo.__args__  # type: ignore
-        except AttributeError:
-            # Python 3.5.x
-            types = typeinfo.__union_params__  # type: ignore
-
-        for T in types:
+        for T in union_types(typeinfo):
             try:
                 check_option_type(name, value, T)
             except TypeError:
@@ -31,12 +58,7 @@ def check_option_type(name: str, value: typing.Any, typeinfo: typing.Any) -> Non
                 return
         raise e
     elif typename.startswith("typing.Tuple"):
-        try:
-            types = typeinfo.__args__  # type: ignore
-        except AttributeError:
-            # Python 3.5.x
-            types = typeinfo.__tuple_params__  # type: ignore
-
+        types = tuple_types(typeinfo)
         if not isinstance(value, (tuple, list)):
             raise e
         if len(types) != len(value):
@@ -45,11 +67,7 @@ def check_option_type(name: str, value: typing.Any, typeinfo: typing.Any) -> Non
             check_option_type("{}[{}]".format(name, i), x, T)
         return
     elif typename.startswith("typing.Sequence"):
-        try:
-            T = typeinfo.__args__[0]  # type: ignore
-        except AttributeError:
-            # Python 3.5.0
-            T = typeinfo.__parameters__[0]  # type: ignore
+        T = sequence_type(typeinfo)
         if not isinstance(value, (tuple, list)):
             raise e
         for v in value:

setup.cfg

@@ -75,7 +75,6 @@ exclude =
     mitmproxy/proxy/protocol/tls.py
     mitmproxy/proxy/root_context.py
     mitmproxy/proxy/server.py
-    mitmproxy/stateobject.py
     mitmproxy/utils/bits.py
     pathod/language/actions.py
     pathod/language/base.py