mitmproxy/mitmproxy/stateobject.py
Aldo Cortesi 5192810ff6 Make mypy succeed with imports on master.py
We get little benefit from our mypy QA checks at the moment, because we skip
imports. This patch is what's needed to make mypy succeed with imports on a
single file: master.py

It also updates mypy to the current version, and enables a QA check.

Mypy bugs I encountered:

dict.update with kwargs not supported:

https://github.com/python/mypy/issues/1031

property setters and getters must be adjacent:

https://github.com/python/mypy/issues/1465
2017-03-17 08:13:47 +13:00

78 lines
2.6 KiB
Python

from typing import Any
from typing import List
from typing import MutableMapping # noqa
from mitmproxy.types 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
class StateObject(serializable.Serializable):
"""
An object with serializable state.
State attributes can either be serializable types(str, tuple, bool, ...)
or StateObject instances themselves.
"""
_stateobject_attributes = None # type: MutableMapping[str, Any]
"""
An attribute-name -> class-or-type dict containing all attributes that
should be serialized. If the attribute is a class, it must implement the
Serializable protocol.
"""
def get_state(self):
"""
Retrieve object state.
"""
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
return state
def set_state(self, state):
"""
Load object state from data returned by a get_state call.
"""
state = state.copy()
for attr, cls in self._stateobject_attributes.items():
val = state.pop(attr)
if val is None:
setattr(self, attr, val)
else:
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))
if state:
raise RuntimeWarning("Unexpected State in __setstate__: {}".format(state))