Merge pull request #921 from mitmproxy/model-cleanup

Model Cleanup

1 files changed, 11 insertions, 8 deletions

diff --git a/libmproxy/models/flow.py b/libmproxy/models/flow.py
index b4e8cb88..10255dad 100644
--- a/libmproxy/models/flow.py
+++ b/libmproxy/models/flow.py
@@ -45,7 +45,7 @@ class Error(stateobject.StateObject):
         # the default implementation assumes an empty constructor. Override
         # accordingly.
         f = cls(None)
-        f.load_state(state)
+        f.set_state(state)
         return f
 
     def copy(self):
@@ -86,16 +86,19 @@ class Flow(stateobject.StateObject):
         intercepted=bool
     )
 
-    def get_state(self, short=False):
-        d = super(Flow, self).get_state(short)
+    def get_state(self):
+        d = super(Flow, self).get_state()
         d.update(version=version.IVERSION)
         if self._backup and self._backup != d:
-            if short:
-                d.update(modified=True)
-            else:
-                d.update(backup=self._backup)
+            d.update(backup=self._backup)
         return d
 
+    def set_state(self, state):
+        state.pop("version")
+        if "backup" in state:
+            self._backup = state.pop("backup")
+        super(Flow, self).set_state(state)
+
     def __eq__(self, other):
         return self is other
 
@@ -133,7 +136,7 @@ class Flow(stateobject.StateObject):
             Revert to the last backed up state.
         """
         if self._backup:
-            self.load_state(self._backup)
+            self.set_state(self._backup)
             self._backup = None
 
     def kill(self, master):