fix portfolio set_state

ricequant · Mar 28, 2017 · ad25758 · ad25758
1 parent fd333e6
commit ad25758
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/rqalpha/model/portfolio.py b/rqalpha/model/portfolio.py
@@ -59,7 +59,12 @@ def set_state(self, state):
         self._static_unit_net_value = value['static_unit_net_value']
         self._units = value['units']
         for k, v in six.iteritems(value['accounts']):
-            self._accounts[k].set_state(v)
+            if k == 'ACCOUNT_TYPE.STOCK':
+                self._accounts[ACCOUNT_TYPE.STOCK].set_state(v)
+            elif k == 'ACCOUNT_TYPE.FUTURE':
+                self._accounts[ACCOUNT_TYPE.FUTURE].set_state(v)
+            else:
+                raise NotImplementedError
 
     def _post_settlement(self, event):
         self._static_unit_net_value = self.unit_net_value