From 1b2c2c34b6b883ca6125fa5d09fd05a38f49108c Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 22 Sep 2012 01:46:53 +0100
Subject: [PATCH] Make CompositeCaster easier to subclass

---
 lib/extras.py              | 26 +++++++++++++++++++-------
 tests/test_types_extras.py | 27 ++++++++++++++++++++++++++-
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/lib/extras.py b/lib/extras.py
index 3fd1fb0c..90652f4c 100644
--- a/lib/extras.py
+++ b/lib/extras.py
@@ -854,8 +854,13 @@ class CompositeCaster(object):
                 "expecting %d components for the type %s, %d found instead" %
                 (len(self.atttypes), self.name, len(tokens)))
 
-        return self._ctor(curs.cast(oid, token)
-            for oid, token in zip(self.atttypes, tokens))
+        values = [ curs.cast(oid, token)
+            for oid, token in zip(self.atttypes, tokens) ]
+
+        return self.make(values)
+
+    def make(self, values):
+        return self._ctor(values)
 
     _re_tokenize = regex.compile(r"""
   \(? ([,)])                        # an empty token, representing NULL
@@ -937,10 +942,10 @@ ORDER BY attnum;
         array_oid = recs[0][1]
         type_attrs = [ (r[2], r[3]) for r in recs ]
 
-        return CompositeCaster(tname, type_oid, type_attrs,
+        return self(tname, type_oid, type_attrs,
             array_oid=array_oid)
 
-def register_composite(name, conn_or_curs, globally=False):
+def register_composite(name, conn_or_curs, globally=False, factory=None):
     """Register a typecaster to convert a composite type into a tuple.
 
     :param name: the name of a PostgreSQL composite type, e.g. created using
@@ -950,14 +955,21 @@ def register_composite(name, conn_or_curs, globally=False):
         object, unless *globally* is set to `!True`
     :param globally: if `!False` (default) register the typecaster only on
         *conn_or_curs*, otherwise register it globally
-    :return: the registered `CompositeCaster` instance responsible for the
-        conversion
+    :param factory: if specified it should be a `CompositeCaster` subclass: use
+        it to :ref:`customize how to cast composite types <custom-composite>`
+    :return: the registered `CompositeCaster` or *factory* instance
+        responsible for the conversion
 
     .. versionchanged:: 2.4.3
         added support for array of composite types
+    .. versionchanged:: 2.4.6
+        added the *factory* parameter
 
     """
-    caster = CompositeCaster._from_db(name, conn_or_curs)
+    if factory is None:
+        factory = CompositeCaster
+
+    caster = factory._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:
diff --git a/tests/test_types_extras.py b/tests/test_types_extras.py
index 49be3908..b8f72419 100755
--- a/tests/test_types_extras.py
+++ b/tests/test_types_extras.py
@@ -736,7 +736,7 @@ class AdaptTypeTestCase(unittest.TestCase):
         self.assertEqual(r[0], (2, 'test2'))
         self.assertEqual(r[1], [(3, 'testc', 2), (4, 'testd', 2)])
 
-    @skip_if_no_hstore
+    @skip_if_no_composite
     def test_non_dbapi_connection(self):
         from psycopg2.extras import RealDictConnection
         from psycopg2.extras import register_composite
@@ -760,6 +760,31 @@ class AdaptTypeTestCase(unittest.TestCase):
         finally:
             conn.close()
 
+    @skip_if_no_composite
+    def test_subclass(self):
+        oid = self._create_type("type_isd",
+            [('anint', 'integer'), ('astring', 'text'), ('adate', 'date')])
+
+        from psycopg2.extras import register_composite, CompositeCaster
+
+        class DictComposite(CompositeCaster):
+            def make(self, values):
+                return dict(zip(self.attnames, values))
+
+        t = register_composite('type_isd', self.conn, factory=DictComposite)
+
+        self.assertEqual(t.name, 'type_isd')
+        self.assertEqual(t.oid, oid)
+
+        curs = self.conn.cursor()
+        r = (10, 'hello', date(2011,1,2))
+        curs.execute("select %s::type_isd;", (r,))
+        v = curs.fetchone()[0]
+        self.assert_(isinstance(v, dict))
+        self.assertEqual(v['anint'], 10)
+        self.assertEqual(v['astring'], "hello")
+        self.assertEqual(v['adate'], date(2011,1,2))
+
     def _create_type(self, name, fields):
         curs = self.conn.cursor()
         try: