From 8963b8adcb5a70e25b3a48ccf88ece5dd8ee44f7 Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 22 Sep 2011 19:55:33 +0200
Subject: [PATCH] Added support for arrays of composite types

---
 NEWS                  |  2 +-
 lib/extras.py         | 37 +++++++++++++++++++++++++++++++------
 tests/types_extras.py | 23 +++++++++++++++++++++++
 3 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/NEWS b/NEWS
index cec9a0d9..bb03987a 100644
--- a/NEWS
+++ b/NEWS
@@ -3,7 +3,7 @@ What's new in psycopg 2.4.3
 
   - Added 'new_array_type()' function for easy creation of array
     typecasters.
-  - Added support for arrays of hstores (ticket #66).
+  - Added support for arrays of hstores and composite types (ticket #66).
   - Fixed segfault in case of transaction started with connection lost
     (and possibly other events).
   - Rollback connections in transaction or in error before putting them
diff --git a/lib/extras.py b/lib/extras.py
index 0ef44a0a..78f962eb 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -778,6 +778,9 @@ class CompositeCaster(object):
     """Helps conversion of a PostgreSQL composite type into a Python object.
 
     The class is usually created by the `register_composite()` function.
+    You may want to create and register manually instances of the class if
+    querying the database at registration time is not desirable (such as when
+    using an :ref:`asynchronous connections <async-support>`).
 
     .. attribute:: name
 
@@ -787,6 +790,10 @@ class CompositeCaster(object):
 
         The oid of the PostgreSQL type.
 
+    .. attribute:: array_oid
+
+        The oid of the PostgreSQL array type, if available.
+
     .. attribute:: type
 
         The type of the Python objects returned. If :py:func:`collections.namedtuple()`
@@ -802,14 +809,20 @@ class CompositeCaster(object):
         List of component type oids of the type to be casted.
 
     """
-    def __init__(self, name, oid, attrs):
+    def __init__(self, name, oid, attrs, array_oid=None):
         self.name = name
         self.oid = oid
+        self.array_oid = array_oid
 
         self.attnames = [ a[0] for a in attrs ]
         self.atttypes = [ a[1] for a in attrs ]
         self._create_type(name, self.attnames)
         self.typecaster = _ext.new_type((oid,), name, self.parse)
+        if array_oid:
+            self.array_typecaster = _ext.new_array_type(
+                (array_oid,), "%sARRAY" % name, self.typecaster)
+        else:
+            self.array_typecaster = None
 
     def parse(self, s, curs):
         if s is None:
@@ -881,15 +894,18 @@ class CompositeCaster(object):
             tname = name
             schema = 'public'
 
+        # column typarray not available before PG 8.3
+        typarray = conn.server_version >= 80300 and "typarray" or "NULL"
+
         # get the type oid and attributes
         curs.execute("""\
-SELECT t.oid, attname, atttypid
+SELECT t.oid, %s, attname, atttypid
 FROM pg_type t
 JOIN pg_namespace ns ON typnamespace = ns.oid
 JOIN pg_attribute a ON attrelid = typrelid
-WHERE typname = %s and nspname = %s
+WHERE typname = %%s and nspname = %%s
 ORDER BY attnum;
-""", (tname, schema))
+""" % typarray, (tname, schema))
 
         recs = curs.fetchall()
 
@@ -903,9 +919,11 @@ ORDER BY attnum;
                 "PostgreSQL type '%s' not found" % name)
 
         type_oid = recs[0][0]
-        type_attrs = [ (r[1], r[2]) for r in recs ]
+        array_oid = recs[0][1]
+        type_attrs = [ (r[2], r[3]) for r in recs ]
 
-        return CompositeCaster(tname, type_oid, type_attrs)
+        return CompositeCaster(tname, type_oid, type_attrs,
+            array_oid=array_oid)
 
 def register_composite(name, conn_or_curs, globally=False):
     """Register a typecaster to convert a composite type into a tuple.
@@ -919,10 +937,17 @@ def register_composite(name, conn_or_curs, globally=False):
         *conn_or_curs*, otherwise register it globally
     :return: the registered `CompositeCaster` instance responsible for the
         conversion
+
+    .. versionchanged:: 2.4.3
+        added support for array of composite types
+
     """
     caster = CompositeCaster._from_db(name, conn_or_curs)
     _ext.register_type(caster.typecaster, not globally and conn_or_curs or None)
 
+    if caster.array_typecaster is not None:
+        _ext.register_type(caster.array_typecaster, not globally and conn_or_curs or None)
+
     return caster
 
 
diff --git a/tests/types_extras.py b/tests/types_extras.py
index 7c0ca910..c95a18c6 100755
--- a/tests/types_extras.py
+++ b/tests/types_extras.py
@@ -620,6 +620,29 @@ class AdaptTypeTestCase(unittest.TestCase):
         curs.execute("select (4,8)::typens.typens_ii")
         self.assertEqual(curs.fetchone()[0], (4,8))
 
+    @skip_if_no_composite
+    @skip_before_postgres(8, 3)
+    def test_composite_array(self):
+        oid = self._create_type("type_isd",
+            [('anint', 'integer'), ('astring', 'text'), ('adate', 'date')])
+
+        t = psycopg2.extras.register_composite("type_isd", self.conn)
+
+        curs = self.conn.cursor()
+        r1 = (10, 'hello', date(2011,1,2))
+        r2 = (20, 'world', date(2011,1,3))
+        curs.execute("select %s::type_isd[];", ([r1, r2],))
+        v = curs.fetchone()[0]
+        self.assertEqual(len(v), 2)
+        self.assert_(isinstance(v[0], t.type))
+        self.assertEqual(v[0][0], 10)
+        self.assertEqual(v[0][1], "hello")
+        self.assertEqual(v[0][2], date(2011,1,2))
+        self.assert_(isinstance(v[1], t.type))
+        self.assertEqual(v[1][0], 20)
+        self.assertEqual(v[1][1], "world")
+        self.assertEqual(v[1][2], date(2011,1,3))
+
     def _create_type(self, name, fields):
         curs = self.conn.cursor()
         try: