More XML fixes..

James William Pye · James William Pye · commit b72b88cb7faf · 2010-09-03T10:51:04.000-07:00
diff --git a/postgresql/test/test_driver.py b/postgresql/test/test_driver.py
@@ -1130,8 +1130,7 @@ def testDomainSupport(self):
 					)
 				)
 
-	@pg_tmp
-	def testXML(self):
+	def check_xml(self):
 		try:
 			xml = db.prepare('select $1::xml')
 			textxml = db.prepare('select $1::text::xml')
@@ -1184,6 +1183,23 @@ def tostr(x):
 			(tostr(foo), tostr(bar))
 		)
 
+	@pg_tmp
+	def testXML(self):
+		self.check_xml()
+
+	@pg_tmp
+	def testXML_ascii(self):
+		# check a non-utf8 encoding (3.2 and up)
+		db.settings['client_encoding'] = 'sql_ascii'
+		self.check_xml()
+
+	@pg_tmp
+	def testXML_utf8(self):
+		# in 3.2 we always serialize at utf-8, so check that
+		# that path is being ran by forcing the client_encoding to utf8.
+		db.settings['client_encoding'] = 'utf8'
+		self.check_xml()
+
 	@pg_tmp
 	def testUUID(self):
 		# doesn't exist in all versions supported by py-postgresql.
diff --git a/postgresql/types/io/stdlib_xml_etree.py b/postgresql/types/io/stdlib_xml_etree.py
@@ -42,15 +42,15 @@ def xml_io_factory(typoid, typio, c = compose):
 		)
 else:
 	# New etree tostring API.
-	def xml_pack(xml, encoding,
+	def xml_pack(xml, encoding, encoder,
 		tostr = etree.tostring, et = etree.ElementTree,
 		str = str, isinstance = isinstance, tuple = tuple,
 	):
 		if isinstance(xml, bytes):
 			return xml
 		if isinstance(xml, str):
 			# If it's a string, encode and return.
-			return xml.encode(encoding)
+			return encoder(xml)
 		elif isinstance(xml, tuple):
 			# If it's a tuple, encode and return the joined items.
 			# We do not accept lists here--emphasizing lists being used for ARRAY
@@ -59,7 +59,7 @@ def xml_pack(xml, encoding,
 			# 3.2
 			# XXX: tostring doesn't include declaration with utf-8?
 			x = b''.join(
-				x.encode(encoding) if isinstance(x, str) else
+				x.encode('utf-8') if isinstance(x, str) else
 				tostr(x, encoding = "utf-8")
 				for x in xml
 			)
@@ -71,11 +71,11 @@ def xml_pack(xml, encoding,
 		if encoding in ('utf8','utf-8'):
 			return x
 		else:
-			return x.decode('utf-8').encode(encoding)
+			return encoder(x.decode('utf-8'))
 
 	def xml_io_factory(typoid, typio, c = compose):
-		def local_xml_pack(x, typio = typio, xml_pack = xml_pack):
-			return xml_pack(x, typio.encoding)
+		def local_xml_pack(x, encoder = typio.encode, typio = typio, xml_pack = xml_pack):
+			return xml_pack(x, typio.encoding, encoder)
 		return (local_xml_pack, c((typio.decode, xml_unpack)), etree.ElementTree,)
 
 oid_to_io = {
