Implement a sized split generator in python.structlib.

James William Pye · James William Pye · commit 21fe68a9bb23 · 2010-03-02T04:56:26.000-07:00
This will be used to implement hstore.
diff --git a/postgresql/python/structlib.py b/postgresql/python/structlib.py
@@ -1,6 +1,5 @@
 ##
-# copyright 2009, James William Pye
-# http://python.projects.postgresql.org
+# .python.structlib - module for extracting serialized data
 ##
 import struct
 from .functools import Composition as compose
@@ -76,3 +75,30 @@ def mk_pack(x):
 	ushort_pack, ushort_unpack = mk_pack("H")
 	long_pack, long_unpack = mk_pack("l")
 	ulong_pack, ulong_unpack = mk_pack("L")
+
+def split_sized_data(
+	data,
+	ulong_unpack = ulong_unpack,
+	null_field = 0xFFFFFFFF,
+	len = len,
+	errmsg = "insufficient data in field {0}, required {1} bytes, {2} remaining".format
+):
+	"""
+	Given serialized record data, return a tuple of tuples of type Oids and
+	attributes.
+	"""
+	v = memoryview(data)
+	f = 1
+	while v:
+		l = ulong_unpack(v)
+		if l == null_field:
+			v = v[4:]
+			yield None
+			continue
+		l += 4
+		d = v[4:l].tobytes()
+		if len(d) < l-4:
+			raise ValueError(errmsg(f, l - 4, len(d)))
+		v = v[l:]
+		f += 1
+		yield d
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
@@ -850,7 +850,7 @@ def testLoadRows(self):
 			list((x[0] for x in gs.rows())),
 			list(range(1, 10001))
 		)
-		# exercise ``for x in chunks: dst.load(x)``
+		# exercise ``for x in chunks: dst.load_rows(x)``
 		with self.db.connector() as db2:
 			db2.execute(
 				"""
diff --git a/postgresql/test/test_python.py b/postgresql/test/test_python.py
@@ -1,14 +1,16 @@
 ##
-# copyright 2009, James William Pye
-# http://python.projects.postgresql.org
+# .test.test_python
 ##
 import unittest
 import socket
 import errno
+import struct
 from itertools import chain
 from operator import methodcaller
 from contextlib import contextmanager
 
+from ..python.itertools import interlace
+from ..python.structlib import split_sized_data
 from ..python.contextlib import *
 from ..python import functools
 from ..python import itertools
@@ -161,6 +163,27 @@ def testNoCM(self):
 			pass
 		self.failUnlessEqual(foo, None)
 
+def join_sized_data(*data,
+	packL = struct.Struct("!L").pack,
+	getlen = lambda x: len(x) if x is not None else 0xFFFFFFFF
+):
+	return b''.join(interlace(map(packL, map(getlen, data)), (x if x is not None else b'' for x in data)))
+
+class test_structlib(unittest.TestCase):
+	def testSizedSplit(self):
+		sample = [
+			(b'foo', b'bar'),
+			(b'foo', None, b'bar'),
+			(b'foo', None, b'bar'),
+			(b'foo', b'bar'),
+			(),
+			(None,None,None),
+			(b'x', None,None,None, b'yz'),
+		]
+		packed_sample = [join_sized_data(*x) for x in sample]
+		self.failUnlessRaises(ValueError, split_sized_data(b'\xFF\xFF\xFF\x01foo').__next__)
+		self.failUnlessEqual(sample, [tuple(split_sized_data(x)) for x in packed_sample])
+
 if __name__ == '__main__':
 	from types import ModuleType
 	this = ModuleType("this")

Original file line number	Diff line number	Diff line change
`@@ -850,7 +850,7 @@ def testLoadRows(self):`
`850`	`850`	`list((x[0] for x in gs.rows())),`
`851`	`851`	`list(range(1, 10001))`
`852`	`852`	`)`
`853`		- # exercise ``for x in chunks: dst.load(x)``
	`853`	+ # exercise ``for x in chunks: dst.load_rows(x)``
`854`	`854`	`with self.db.connector() as db2:`
`855`	`855`	`db2.execute(`
`856`	`856`	`"""`