~ubuntu-branches/ubuntu/precise/pyzmq/precise

« back to all changes in this revision

Viewing changes to examples/serialization/serialsocket.py

  • Committer: Package Import Robot
  • Author(s): Debian Python Modules Team
  • Date: 2011-09-23 00:16:39 UTC
  • mfrom: (1.1.4 upstream)
  • Revision ID: package-import@ubuntu.com-20110923001639-girjqodpb7uv17yu
Tags: 2.1.9-1
* New upstream version
  - should build on kFreeBSD without patches (Closes: #637777).
* Build-depend on zeromq 2.1.9

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
"""A Socket subclass that adds some serialization methods."""
 
2
 
 
3
import zlib
 
4
import cPickle as pickle
 
5
 
 
6
import numpy
 
7
 
 
8
import zmq
 
9
 
 
10
class SerializingSocket(zmq.Socket):
 
11
    """A class with some extra serialization methods
 
12
    
 
13
    send_zipped_pickle is just like send_pyobj, but uses
 
14
    zlib to compress the stream before sending.
 
15
    
 
16
    send_array sends numpy arrays with metadata necessary
 
17
    for reconstructing the array on the other side (dtype,shape).
 
18
    """
 
19
    
 
20
    def send_zipped_pickle(self, obj, flags=0, protocol=-1):
 
21
        """pack and compress an object with pickle and zlib."""
 
22
        pobj = pickle.dumps(obj, protocol)
 
23
        zobj = zlib.compress(pobj)
 
24
        print 'zipped pickle is %i bytes'%len(zobj)
 
25
        return self.send(zobj, flags=flags)
 
26
    
 
27
    def recv_zipped_pickle(self, flags=0):
 
28
        """reconstruct a Python object sent with zipped_pickle"""
 
29
        zobj = self.recv(flags)
 
30
        pobj = zlib.decompress(zobj)
 
31
        return pickle.loads(pobj)
 
32
 
 
33
    def send_array(self, A, flags=0, copy=True, track=False):
 
34
        """send a numpy array with metadata"""
 
35
        md = dict(
 
36
            dtype = str(A.dtype),
 
37
            shape = A.shape,
 
38
        )
 
39
        self.send_json(md, flags|zmq.SNDMORE)
 
40
        return self.send(A, flags, copy=copy, track=track)
 
41
 
 
42
    def recv_array(self, flags=0, copy=True, track=False):
 
43
        """recv a numpy array"""
 
44
        md = self.recv_json(flags=flags)
 
45
        msg = self.recv(flags=flags, copy=copy, track=track)
 
46
        buf = buffer(msg)
 
47
        A = numpy.frombuffer(buf, dtype=md['dtype'])
 
48
        return A.reshape(md['shape'])
 
49
    
 
50
 
 
51
if __name__ == '__main__':
 
52
    ctx = zmq.Context.instance()
 
53
    req = SerializingSocket(ctx, zmq.REQ)
 
54
    rep = SerializingSocket(ctx, zmq.REP)
 
55
    
 
56
    rep.bind('inproc://a')
 
57
    req.connect('inproc://a')
 
58
    A = numpy.ones((1024,1024))
 
59
    print "Array is %i bytes"%len(buffer(A))
 
60
    
 
61
    # send/recv with pickle+zip
 
62
    req.send_zipped_pickle(A)
 
63
    B = rep.recv_zipped_pickle()
 
64
    # now try non-copying version
 
65
    rep.send_array(A, copy=False)
 
66
    C = req.recv_array(copy=False)
 
67
    print ("Checking zipped pickle...")
 
68
    print ("Okay" if (A==B).all() else "Failed")
 
69
    print ("Checking send_array...")
 
70
    print ("Okay" if (C==B).all() else "Failed")