Merge pull request #120 from mitmproxy/model-cleanup

Model Cleanup

1 files changed, 29 insertions, 3 deletions

diff --git a/netlib/http/message.py b/netlib/http/message.py
index 28f55fa2..e3d8ce37 100644
--- a/netlib/http/message.py
+++ b/netlib/http/message.py
@@ -4,9 +4,9 @@ import warnings
 
 import six
 
+from .headers import Headers
 from .. import encoding, utils
 
-
 CONTENT_MISSING = 0
 
 if six.PY2:  # pragma: nocover
@@ -18,7 +18,7 @@ else:
     _always_bytes = lambda x: utils.always_bytes(x, "utf-8", "surrogateescape")
 
 
-class MessageData(object):
+class MessageData(utils.Serializable):
     def __eq__(self, other):
         if isinstance(other, MessageData):
             return self.__dict__ == other.__dict__
@@ -27,8 +27,24 @@ class MessageData(object):
     def __ne__(self, other):
         return not self.__eq__(other)
 
+    def set_state(self, state):
+        for k, v in state.items():
+            if k == "headers":
+                v = Headers.from_state(v)
+            setattr(self, k, v)
+
+    def get_state(self):
+        state = vars(self).copy()
+        state["headers"] = state["headers"].get_state()
+        return state
+
+    @classmethod
+    def from_state(cls, state):
+        state["headers"] = Headers.from_state(state["headers"])
+        return cls(**state)
+
 
-class Message(object):
+class Message(utils.Serializable):
     def __init__(self, data):
         self.data = data
 
@@ -40,6 +56,16 @@ class Message(object):
     def __ne__(self, other):
         return not self.__eq__(other)
 
+    def get_state(self):
+        return self.data.get_state()
+
+    def set_state(self, state):
+        self.data.set_state(state)
+
+    @classmethod
+    def from_state(cls, state):
+        return cls(**state)
+
     @property
     def headers(self):
         """