Implement state loading that doesn't change object identity.

We need this to let us load state from copied Flows returned from scripts.

1 files changed, 18 insertions, 3 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 03b8b309..3520cc93 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -99,11 +99,26 @@ class Flow:
 
     def load_state(self, state):
         self._backup = state["backup"]
-        self.request = proxy.Request.from_state(state["request"])
+        if self.request:
+            self.request.load_state(state["request"])
+        else:
+            self.request = proxy.Request.from_state(state["request"])
+
         if state["response"]:
-            self.response = proxy.Response.from_state(self.request, state["response"])
+            if self.response:
+                self.response.load_state(state["response"])
+            else:
+                self.response = proxy.Response.from_state(self.request, state["response"])
+        else:
+            self.response = None
+
         if state["error"]:
-            self.error = proxy.Error.from_state(state["error"])
+            if self.error:
+                self.error.load_state(state["error"])
+            else:
+                self.error = proxy.Error.from_state(state["error"])
+        else:
+            self.error = None
 
     @classmethod
     def from_state(klass, state):