polish SimpleStateObject

2 files changed, 42 insertions, 21 deletions

diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 4baee3ee..2caeb011 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -4,6 +4,7 @@
 """
 import hashlib, Cookie, cookielib, copy, re, urlparse, threading
 import time, urllib
+import types
 import tnetstring, filt, script, utils, encoding, proxy
 from email.utils import parsedate_tz, formatdate, mktime_tz
 from netlib import odict, http, certutils, wsgi
@@ -164,41 +165,61 @@ class StateObject:
 class SimpleStateObject(StateObject):
     """
     A StateObject with opionated conventions that tries to keep everything DRY.
+
+    Simply put, you agree on a list of attributes and their type.
+    Attributes can either be primitive types(str, tuple, bool, ...) or StateObject instances themselves.
+    SimpleStateObject uses this information for the default _get_state(), _from_state(s) and _load_state(s) methods.
+    Overriding _get_state or _load_state to add custom adjustments is always possible.
     """
 
-    _stateobject_attributes = None
+    _stateobject_attributes = None  # none by default to raise an exception if definition was forgotten
     """
     An attribute-name -> class-or-type dict containing all attributes that should be serialized
     If the attribute is a class, this class must be a subclass of StateObject.
     """
 
     def _get_state(self):
-        return {attr: (getattr(self, attr)._get_state()
-                       if (type(cls) == 'classobj')
-                       else getattr(self, attr))
+        return {attr: self.__get_state_attr(attr, cls)
                 for attr, cls in self._stateobject_attributes.iteritems()}
 
+    def __get_state_attr(self, attr, cls):
+        """
+        helper for _get_state.
+        returns the value of the given attribute
+        """
+        if getattr(self, attr) is None:
+            return None
+        if isinstance(cls, types.ClassType):
+            return getattr(self, attr)._get_state()
+        else:
+            return getattr(self, attr)
+
     def _load_state(self, state):
         for attr, cls in self._stateobject_attributes.iteritems():
-            self._load_state_attr(attr, cls, state)
-
-    def _load_state_attr(self, attribute, cls, state):
-        if state[attribute] is not None:
-            if type(cls) == 'classobj':
-                assert issubclass(cls, StateObject)
-                curr = getattr(self, attribute)
-                if curr:
-                    curr._load_state(state[attribute])
-                else:
-                    setattr(self, attribute, cls._from_state(state[attribute]))
+            self.__load_state_attr(attr, cls, state)
+
+    def __load_state_attr(self, attr, cls, state):
+        """
+        helper for _load_state.
+        loads the given attribute from the state.
+        """
+        if state[attr] is not None:  # First, catch None as value.
+            if isinstance(cls, types.ClassType):  # Is the attribute a StateObject itself?
+                # FIXME: assertion doesn't hold because of odict at the moment
+                # assert issubclass(cls, StateObject)
+                curr = getattr(self, attr)
+                if curr:  # if the attribute is already present, delegate to the objects ._load_state method.
+                    curr._load_state(state[attr])
+                else: # otherwise, create a new object.
+                    setattr(self, attr, cls._from_state(state[attr]))
             else:
-                setattr(self, attribute, cls(state[attribute]))
+                setattr(self, attr, cls(state[attr]))
         else:
-            setattr(self, attribute, None)
+            setattr(self, attr, None)
 
     @classmethod
     def _from_state(cls, state):
-        f = cls()
+        f = cls()  # the default implementation assumes an empty constructor. Override accordingly.
         f._load_state(state)
         return f
 
diff --git a/libmproxy/protocol/http.py b/libmproxy/protocol/http.py
index 51e73010..4a3933e7 100644
--- a/libmproxy/protocol/http.py
+++ b/libmproxy/protocol/http.py
@@ -58,7 +58,6 @@ class decoded(object):
             self.o.encode(self.ce)
 
 
-
 class BackreferenceMixin(object):
     """
     If an attribute from the _backrefattr tuple is set,
@@ -78,7 +77,7 @@ class BackreferenceMixin(object):
             assert (getattr(value, self._backrefname, self) or self) is self
             setattr(value, self._backrefname, self)
 
-
+# FIXME: Move out of http
 class Error(SimpleStateObject):
     """
         An Error.
@@ -107,7 +106,7 @@ class Error(SimpleStateObject):
         c = copy.copy(self)
         return c
 
-
+# FIXME: Move out of http
 class Flow(SimpleStateObject, BackreferenceMixin):
     def __init__(self, conntype, client_conn, server_conn, error):
         self.conntype = conntype
@@ -167,6 +166,7 @@ class Flow(SimpleStateObject, BackreferenceMixin):
             self._load_state(self._backup)
             self._backup = None
 
+
 class HTTPMessage(SimpleStateObject):
     def __init__(self):
         self.flow = None # Will usually set by backref mixin