[Assorted-commits] SF.net SVN: assorted: [227] python-commons/trunk/src/commons/seqs.py

[<yan...@us...>]

Revision: 227
          http://assorted.svn.sourceforge.net/assorted/?rev=227&view=rev
Author:   yangzhang
Date:     2008-01-12 18:29:58 -0800 (Sat, 12 Jan 2008)

Log Message:
-----------
added safe_pickle, write_pickle; made read_pickle (formerly the AFAIK unused read_objs) more generic

Modified Paths:
--------------
    python-commons/trunk/src/commons/seqs.py

Modified: python-commons/trunk/src/commons/seqs.py
===================================================================
--- python-commons/trunk/src/commons/seqs.py	2008-01-11 04:40:18 UTC (rev 226)
+++ python-commons/trunk/src/commons/seqs.py	2008-01-13 02:29:58 UTC (rev 227)
@@ -1,8 +1,12 @@
 # -*- mode: python; tab-width: 4; indent-tabs-mode: nil; py-indent-offset: 4; -*-
 # vim:ft=python:et:sw=4:ts=4
 
+from __future__ import with_statement
+
 from cStringIO import StringIO
-import cPickle as pickle
+from cPickle import *
+from struct import pack, unpack
+from contextlib import closing
 from itertools import ( chain, count, ifilterfalse, islice,
                         izip, tee )
 
@@ -12,38 +16,52 @@
 @var default_chunk_size: The default chunk size used by L{chunkify}.
 """
 
-default_chunk_size = 4096
+default_chunk_size = 8192
 
-def read_objs( s ):
+def read_pickle( read ):
     """
-    Given an input socket, reads in pickled objects from it. This is a
-    generator which yields those objects as they come. I assume that
-    the pickling format is resistant against partial (incomplete)
-    pickles.
+    Given a reader function L{read}, reads in pickled objects from it. I am a
+    generator which yields unpickled objects. I assume that the pickling
+    is "safe," done using L{safe_pickle}.
 
-    @param stream: The input stream.
-    @type stream: stream
+    @param read: The reader function that reads from a stream. It should take
+    a single argument, the number of bytes to consume.
+    @type read: function
+    """
+    with closing( StringIO() ) as stream:
+        obj = None # return this if we hit eof (not enough bytes read)
 
-    @param chunk_size: The size of the chunk (usually the number of
-    bytes to read).
-    @type chunk_size: int
+        def read_until( target ):
+            remain = target - streamlen( stream )
+            if remain > 0:
+                chunk = read( remain )
+                # append to end
+                stream.seek(0,2)
+                stream.write( chunk )
+            return stream.tell() >= target
+
+        if read_until(4):
+            stream.seek(0)
+            (length,) = unpack('i4', stream.read(4))
+            if read_until(length+4):
+                stream.seek(4)
+                obj = load(stream)
+
+        return ( obj, stream.read() )
+
+def safe_pickle( obj ):
     """
-    stream = StringIO()
-    while True:
-        chunk = s.recv( 8192 )
-        if len( chunk ) == 0: break
-        stream.write( chunk )
-        stream.seek( 0 )
-        while True:
-            try: obj = pickle.load( stream )
-            except ( EOFError, pickle.UnpicklingError, ValueError ): break
-            else:
-                yield obj
-                rem = stream.read()
-                stream.seek( 0 )
-                stream.write( rem )
-                stream.truncate()
+    Pickle L{obj} but prepends the serialized length in bytes.
+    """
+    msg = dumps(obj)
+    return pack('i4',len(msg)) + msg
 
+def write_pickle( obj, write ):
+    """
+    Write L{obj} using function L{write}, in a safe, pickle-able fashion.
+    """
+    return write( safe_pickle( obj ) )
+
 def streamlen( stream ):
     """
     Get the length of a stream (e.g. file stream or StringIO).


This was sent by the SourceForge.net collaborative development platform, the world's largest Open Source development site.