diff --git a/NEWS b/NEWS index 2f3b8650..2f68a3f9 100644 --- a/NEWS +++ b/NEWS @@ -3,6 +3,9 @@ What's new in psycopg 2.4.6 - Added support for backward scrollable cursors. Thanks to Jon Nelson for the initial patch (ticket #108). + - Added a simple way to customize casting of composite types into Python + objects other than namedtuples. Many thanks to Ronan Dunklau and + Tobias Oberstein for the feature development. - connection.reset() implemented using DISCARD ALL on server versions supporting it. - Fixed 'cursor()' arguments propagation in connection subclasses diff --git a/doc/src/extras.rst b/doc/src/extras.rst index ce0e8315..0f76b0dc 100644 --- a/doc/src/extras.rst +++ b/doc/src/extras.rst @@ -198,8 +198,8 @@ after a table row type) into a Python named tuple, or into a regular tuple if >>> cur.fetchone()[0] card(value=8, suit='hearts') -Nested composite types are handled as expected, but the type of the composite -components must be registered as well. +Nested composite types are handled as expected, provided that the type of the +composite components are registered as well. .. doctest:: @@ -214,10 +214,75 @@ components must be registered as well. Adaptation from Python tuples to composite types is automatic instead and requires no adapter registration. + +.. _custom-composite: + +.. Note:: + + If you want to convert PostgreSQL composite types into something different + than a `!namedtuple` you can subclass the `CompositeCaster` overriding + `~CompositeCaster.make()`. For example, if you want to convert your type + into a Python dictionary you can use:: + + >>> class DictComposite(psycopg2.extras.CompositeCaster): + ... def make(self, values): + ... return dict(zip(self.attnames, values)) + + >>> psycopg2.extras.register_composite('card', cur, + ... factory=DictComposite) + + >>> cur.execute("select (8, 'hearts')::card") + >>> cur.fetchone()[0] + {'suit': 'hearts', 'value': 8} + + .. autofunction:: register_composite + .. versionchanged:: 2.4.3 + added support for array of composite types + .. versionchanged:: 2.4.6 + added the *factory* parameter + + .. autoclass:: CompositeCaster + .. automethod:: make + + .. versionadded:: 2.4.6 + + Object attributes: + + .. attribute:: name + + The name of the PostgreSQL type. + + .. attribute:: schema + + The schema where the type is defined. + + .. versionadded:: 2.4.6 + + .. attribute:: oid + + 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()` + is available, it is a named tuple with attributes equal to the type + components. Otherwise it is just the `!tuple` object. + + .. attribute:: attnames + + List of component names of the type to be casted. + + .. attribute:: atttypes + + List of component type oids of the type to be casted. .. index:: diff --git a/lib/extras.py b/lib/extras.py index b01e6e7a..afafd8e2 100644 --- a/lib/extras.py +++ b/lib/extras.py @@ -792,35 +792,10 @@ class CompositeCaster(object): querying the database at registration time is not desirable (such as when using an :ref:`asynchronous connections `). - .. attribute:: name - - The name of the PostgreSQL type. - - .. attribute:: oid - - 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()` - is available, it is a named tuple with attributes equal to the type - components. Otherwise it is just the `!tuple` object. - - .. attribute:: attnames - - List of component names of the type to be casted. - - .. attribute:: atttypes - - List of component type oids of the type to be casted. - """ - def __init__(self, name, oid, attrs, array_oid=None): + def __init__(self, name, oid, attrs, array_oid=None, schema=None): self.name = name + self.schema = schema self.oid = oid self.array_oid = array_oid @@ -844,8 +819,22 @@ 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 a new Python object representing the data being casted. + + *values* is the list of attributes, already casted into their Python + representation. + + You can subclass this method to :ref:`customize the composite cast + `. + """ + + return self._ctor(values) _re_tokenize = regex.compile(r""" \(? ([,)]) # an empty token, representing NULL @@ -927,10 +916,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, - array_oid=array_oid) + return self(tname, type_oid, type_attrs, + array_oid=array_oid, schema=schema) -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 @@ -940,14 +929,15 @@ 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 - - .. versionchanged:: 2.4.3 - added support for array of composite types - + :param factory: if specified it should be a `CompositeCaster` subclass: use + it to :ref:`customize how to cast composite types ` + :return: the registered `CompositeCaster` or *factory* instance + responsible for the conversion """ - 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..f5a60818 100755 --- a/tests/test_types_extras.py +++ b/tests/test_types_extras.py @@ -532,6 +532,7 @@ class AdaptTypeTestCase(unittest.TestCase): t = psycopg2.extras.register_composite("type_isd", self.conn) self.assertEqual(t.name, 'type_isd') + self.assertEqual(t.schema, 'public') self.assertEqual(t.oid, oid) self.assert_(issubclass(t.type, tuple)) self.assertEqual(t.attnames, ['anint', 'astring', 'adate']) @@ -655,6 +656,7 @@ class AdaptTypeTestCase(unittest.TestCase): [("a", "integer"), ("b", "integer")]) t = psycopg2.extras.register_composite( "typens.typens_ii", self.conn) + self.assertEqual(t.schema, 'typens') curs.execute("select (4,8)::typens.typens_ii") self.assertEqual(curs.fetchone()[0], (4,8)) @@ -736,7 +738,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 +762,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: