Python unpickle a object with a class instance inside

Question

I am using python requests library and trying to persist the session.

Since I have multiple IPs on my host, I created the following method in order to make session bind to a specific IP.

class SourceAddressAdapter(HTTPAdapter):
    def __init__(self, source_address, **kwargs):
        self.source_address = source_address
        super(SourceAddressAdapter, self).__init__(**kwargs)

    def init_poolmanager(self, connections, maxsize, block=False):
        self.poolmanager = PoolManager(num_pools=connections,
                                   maxsize=maxsize,
                                   block=block,
                                   source_address=self.source_address)

The following code piece is used to invoke this class:

r = requests.Session()
r.mount('http://', SourceAddressAdapter((self.ip,0)))
r.mount('https://', SourceAddressAdapter((self.ip,0)))

After mounting http and https protocol to this adapter, I used pickle to persist the object into redis as follows:

session = pickle.dumps(r)
redis.hset('sessions',id,session)

The problem occurred when I tried to unpickle the session object:

s=redis.hget('sessions', id)
pickle.loads(s)


Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python2.7/pickle.py", line 1382, in loads
    return Unpickler(file).load()
  File "/usr/lib/python2.7/pickle.py", line 858, in load
    dispatch[key](self)
  File "/usr/lib/python2.7/pickle.py", line 1217, in load_build
    setstate(state)
  File "/usr/local/lib/python2.7/dist-packages/requests/adapters.py", line 114, in __setstate__
    block=self._pool_block)
  File "network_driver.py", line 158, in init_poolmanager
    source_address=self.source_address)
AttributeError: 'SourceAddressAdapter' object has no attribute 'source_address'

It complained the SourceAddressAdapter does not have source_address attribute. Before I added this class SourceAddressAdapter to my session, the serialization worked well.

So I guess this is a problem with customized class pickling/unpickling.

---

UPDATE:

It works after I added __getstate__ and __setstate__ method into SourceAddressAdapter

def __getstate__(self):
    # it calls HTTPAdapter's __getstate__()
    state = super(SourceAddressAdapter, self).__getstate__() 
    state['source_address'] = self.source_address
    return state

def __setstate__(self,state):
    self.source_address = state['source_address']
    # Call HTTPAdapter's __setstate__ function to pack the attributes in parent class 
    super(SourceAddressAdapter, self).__setstate__(state)

Answer 1

I believe the problem is that the HTTPAdapter class defines a __setstate__ method. This function is called upon unpickling, and restores the instance to the pickled state. However, the HTTPAdapter knows nothing of your source_address attribute, so that attribute isn't restored (or maybe not even pickled in the first place).

To fix this, you'll need to override the __setstate__ function, somewhat like this:

def __setstate__(self, state):
    self.source_address= state['source_address'] # do this before calling __setstate__
    HTTPAdapter.__setstate__(self, state)

And, as previously mentioned, you might also have to override the __getstate__ function so that the source_address gets pickled.