Add XML tests and expand functionality.

James William Pye · James William Pye · commit 4100d606c00b · 2009-03-23T19:11:01.000-07:00
Use a tuple to identify a sequence of fragments instead of a list, this allows
for unambiguous usage with arrays.

Also, if given text, just submit the encoded text, don't bother parsing and
serializing it.. (still not sure about this one, but it allows the user to
submit fragments).
diff --git a/postgresql/protocol/typio.py b/postgresql/protocol/typio.py
@@ -620,15 +620,25 @@ def anyarray_unpack(self, adata):
 		return anyarray_unpack(self.resolve_unpack, adata)
 
 	def xml_pack(self, xml):
-		return self._encode(etree.tostring(xml))[0]
+		if isinstance(xml, str):
+			# if it's a string, encode and return.
+			return self._encode(xml)[0]
+		elif isinstance(xml, tuple):
+			# if it's a tuple, encode and return the joined items.
+			return b''.join((
+				self._encode(
+					x if isinstance(x, str) else pg_types.etree.tostring(x)
+				)[0] for x in xml
+			))
+		return self._encode(pg_types.etree.tostring(xml))[0]
 
 	def xml_unpack(self, xmldata):
 		xml_or_frag = self._decode(xmldata)[0]
 		try:
 			return pg_types.etree.XML(xml_or_frag)
 		except Exception:
 			# try it again, but return the sequence of children.
-			return list(pg_types.etree.XML('<x>' + xml_or_frag + '</x>'))
+			return tuple(pg_types.etree.XML('<x>' + xml_or_frag + '</x>'))
 
 	def attribute_map(self, pq_descriptor):
 		return zip(self.decodes(pq_descriptor.keys()), count())
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
@@ -672,6 +672,52 @@ def testTypes(self):
 					)
 				)
 
+	def testXML(self):
+		try:
+			xml = self.db.prepare('select $1::xml')
+			textxml = self.db.prepare('select $1::text::xml')
+			r = textxml.first('<foo/>')
+		except (pg_exc.FeatureError, pg_exc.UndefinedObjectError):
+			return
+		foo = pg_types.etree.XML('<foo/>')
+		bar = pg_types.etree.XML('<bar/>')
+		tostr = pg_types.etree.tostring
+		self.failUnlessEqual(tostr(xml.first(foo)), tostr(foo))
+		self.failUnlessEqual(tostr(xml.first(bar)), tostr(bar))
+		self.failUnlessEqual(tostr(textxml.first('<foo/>')), tostr(foo))
+		self.failUnlessEqual(tostr(textxml.first('<foo/>')), tostr(foo))
+		self.failUnlessEqual(tostr(xml.first(pg_types.etree.XML('<foo/>'))), tostr(foo))
+		self.failUnlessEqual(tostr(textxml.first('<foo/>')), tostr(foo))
+		# test fragments
+		self.failUnlessEqual(
+			tuple(
+				tostr(x) for x in xml.first('<foo/><bar/>')
+			), (tostr(foo), tostr(bar))
+		)
+		self.failUnlessEqual(
+			tuple(
+				tostr(x) for x in textxml.first('<foo/><bar/>')
+			),
+			(tostr(foo), tostr(bar))
+		)
+		# mixed text and etree.
+		self.failUnlessEqual(
+			tuple(
+				tostr(x) for x in xml.first((
+					'<foo/>', bar,
+				))
+			),
+			(tostr(foo), tostr(bar))
+		)
+		self.failUnlessEqual(
+			tuple(
+				tostr(x) for x in xml.first((
+					'<foo/>', bar,
+				))
+			),
+			(tostr(foo), tostr(bar))
+		)
+
 	def testSyntaxError(self):
 		self.failUnlessRaises(
 			pg_exc.SyntaxError,
