From 390a9c2451cd2906ef5ff9cd3bb75074fb260e6e Mon Sep 17 00:00:00 2001
From: Federico Di Gregorio <fog@initd.org>
Date: Sun, 4 Oct 2009 12:34:02 +0200
Subject: [PATCH] Added support for UUID arrays

---
 ChangeLog             |  2 ++
 lib/extras.py         | 27 ++++++++++++++++++++++++---
 tests/types_extras.py | 22 +++++++++++++++++++++-
 3 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/ChangeLog b/ChangeLog
index 8153d2c5..347fc365 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,5 +1,7 @@
 2009-10-04  Federico Di Gregorio  <fog@initd.org>
 
+	* lib/extras.py: added support for UUID arrays.
+
 	* psycopg/connection_int.c: applied patch from Richard Davies to avoid
 	deadlocks with multiple threads using the same connection.
 
diff --git a/lib/extras.py b/lib/extras.py
index e4665701..d0b25f20 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -340,13 +340,34 @@ try:
             
         __str__ = getquoted
 
-    def register_uuid(oid=None, conn_or_curs=None):
+    def register_uuid(oids=None, conn_or_curs=None):
         """Create the UUID type and an uuid.UUID adapter."""
-        if not oid: oid = 2950
-        _ext.UUID = _ext.new_type((oid, ), "UUID",
+        if not oids:
+            oid1 = 2950
+            oid2 = 2951
+        elif type(oids) == list:
+            oid1, oid2 = oids
+        else:
+            oid1 = oids
+            oid2 = 2951
+
+        def parseUUIDARRAY(data, cursor):
+            if data is None:
+                return None
+            elif data == '{}':
+                return []
+            else:
+                return [((len(x) > 0 and x != 'NULL') and uuid.UUID(x) or None)
+                        for x in data[1:-1].split(',')]
+
+        _ext.UUID = _ext.new_type((oid1, ), "UUID",
                 lambda data, cursor: data and uuid.UUID(data) or None)
+        _ext.UUIDARRAY = _ext.new_type((oid2,), "UUID[]", parseUUIDARRAY)
+
         _ext.register_type(_ext.UUID, conn_or_curs)
+        _ext.register_type(_ext.UUIDARRAY, conn_or_curs)
         _ext.register_adapter(uuid.UUID, UUID_adapter)
+
         return _ext.UUID
 
 except ImportError, e:
diff --git a/tests/types_extras.py b/tests/types_extras.py
index 1342dd15..6cd55c1c 100644
--- a/tests/types_extras.py
+++ b/tests/types_extras.py
@@ -42,13 +42,33 @@ class TypesExtrasTests(unittest.TestCase):
             psycopg2.extras.register_uuid()
         except:
             return
-        u = uuid.UUID('9c6d5a77-7256-457e-9461-347b4358e350');
+        u = uuid.UUID('9c6d5a77-7256-457e-9461-347b4358e350')
         s = self.execute("SELECT %s AS foo", (u,))
         self.failUnless(u == s)
         # must survive NULL cast to a uuid
         s = self.execute("SELECT NULL::uuid AS foo")
         self.failUnless(s is None)
 
+    def testUUIDARRAY(self):
+        try:
+            import uuid
+            psycopg2.extras.register_uuid()
+        except:
+            return
+        u = [uuid.UUID('9c6d5a77-7256-457e-9461-347b4358e350'), uuid.UUID('9c6d5a77-7256-457e-9461-347b4358e352')]
+        s = self.execute("SELECT %s AS foo", (u,))
+        self.failUnless(u == s)
+        # array with a NULL element
+        u = [uuid.UUID('9c6d5a77-7256-457e-9461-347b4358e350'), None]
+        s = self.execute("SELECT %s AS foo", (u,))
+        self.failUnless(u == s)
+        # must survive NULL cast to a uuid[]
+        s = self.execute("SELECT NULL::uuid[] AS foo")
+        self.failUnless(s is None)
+        # what about empty arrays?
+        s = self.execute("SELECT '{}'::uuid[] AS foo")
+        self.failUnless(type(s) == list and len(s) == 0)
+
     def testINET(self):
         psycopg2.extras.register_inet()
         i = "192.168.1.0/24";