psycopg/psycopg2
1
1
Fork 0
mirror of https://github.com/psycopg/psycopg2.git synced 2024-11-23 01:16:34 +03:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'array-typecast' into devel

Browse Source
This commit is contained in:
Daniele Varrazzo 2011-09-22 20:00:08 +02:00
commit 50b445fa12
11 changed files with 229 additions and 21 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
NEWS
View File

@ -1,6 +1,9 @@
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 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
@ -11,6 +14,7 @@ What's new in psycopg 2.4.3
- Fixed interaction between RealDictCursor and named cursors
(ticket #67).
- Dropped limit on the columns length in COPY operations (ticket #68).
- Fixed typecasting of arrays containing consecutive backslashes.
- 'errorcodes' map updated to PostgreSQL 9.1.

3
doc/src/advanced.rst
View File

@ -216,6 +216,9 @@ read:
>>> print type(point), point.x, point.y
<class 'Point'> 10.2 20.3
A typecaster created by `!new_type()` can be also used with
`~psycopg2.extensions.new_array_type()` to create a typecaster converting a
PostgreSQL array into a Python list.
.. index::

19
doc/src/extensions.rst
View File

@ -290,7 +290,7 @@ details.
.. function:: new_type(oids, name, adapter)
Create a new type caster to convert from a PostgreSQL type to a Python
object. The created object must be registered using
object. The object created must be registered using
`register_type()` to be used.
:param oids: tuple of OIDs of the PostgreSQL type to convert.
@ -309,6 +309,23 @@ details.
See :ref:`type-casting-from-sql-to-python` for an usage example.
.. function:: new_array_type(oids, name, base_caster)
Create a new type caster to convert from a PostgreSQL array type to a list
of Python object. The object created must be registered using
`register_type()` to be used.
:param oids: tuple of OIDs of the PostgreSQL type to convert. It should
probably be the oid of the array type (e.g. the ``typarray`` field in
the ``pg_type`` table.
:param name: the name of the new type adapter.
:param base_caster: a Psycopg typecaster, e.g. created using the
`new_type()` function. The caster should be able to parse a single
item of the desired type.
.. versionadded:: 2.4.3
.. function:: register_type(obj [, scope])
Register a type caster created using `new_type()`.

2
lib/extensions.py
View File

@ -57,7 +57,7 @@ except ImportError:
pass
from psycopg2._psycopg import adapt, adapters, encodings, connection, cursor, lobject, Xid
from psycopg2._psycopg import string_types, binary_types, new_type, register_type
from psycopg2._psycopg import string_types, binary_types, new_type, new_array_type, register_type
from psycopg2._psycopg import ISQLQuote, Notify
from psycopg2._psycopg import QueryCanceledError, TransactionRollbackError

67
lib/extras.py
View File

@ -699,7 +699,8 @@ WHERE typname = 'hstore';
return tuple(rv0), tuple(rv1)
def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
def register_hstore(conn_or_curs, globally=False, unicode=False,
oid=None, array_oid=None):
"""Register adapter and typecaster for `!dict`\-\ |hstore| conversions.
:param conn_or_curs: a connection or cursor: the typecaster will be
@ -709,14 +710,18 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
will be `!unicode` instead of `!str`. The option is not available on
Python 3
:param oid: the OID of the |hstore| type if known. If not, it will be
queried on *conn_or_curs*
queried on *conn_or_curs*.
:param array_oid: the OID of the |hstore| array type if known. If not, it
will be queried on *conn_or_curs*.
The connection or cursor passed to the function will be used to query the
database and look for the OID of the |hstore| type (which may be different
across databases). If querying is not desirable (e.g. with
:ref:`asynchronous connections <async-support>`) you may specify it in the
*oid* parameter (it can be found using a query such as :sql:`SELECT
'hstore'::regtype::oid;`).
*oid* parameter, which can be found using a query such as :sql:`SELECT
'hstore'::regtype::oid`. Analogously you can obtain a value for *array_oid*
using a query such as :sql:`SELECT 'hstore[]'::regtype::oid`.
Note that, when passing a dictionary from Python to the database, both
strings and unicode keys and values are supported. Dictionaries returned
@ -730,6 +735,10 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
added the *oid* parameter. If not specified, the typecaster is
installed also if |hstore| is not installed in the :sql:`public`
schema.
.. versionchanged:: 2.4.3
added support for |hstore| array.
"""
if oid is None:
oid = HstoreAdapter.get_oids(conn_or_curs)
@ -738,11 +747,18 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
"hstore type not found in the database. "
"please install it from your 'contrib/hstore.sql' file")
else:
oid = oid[0] # for the moment we don't have a HSTOREARRAY
array_oid = oid[1]
oid = oid[0]
if isinstance(oid, int):
oid = (oid,)
if array_oid is not None:
if isinstance(array_oid, int):
array_oid = (array_oid,)
else:
array_oid = tuple([x for x in array_oid if x])
# create and register the typecaster
if sys.version_info[0] < 3 and unicode:
cast = HstoreAdapter.parse_unicode
@ -753,11 +769,18 @@ def register_hstore(conn_or_curs, globally=False, unicode=False, oid=None):
_ext.register_type(HSTORE, not globally and conn_or_curs or None)
_ext.register_adapter(dict, HstoreAdapter)
if array_oid:
HSTOREARRAY = _ext.new_array_type(array_oid, "HSTOREARRAY", HSTORE)
_ext.register_type(HSTOREARRAY, not globally and conn_or_curs or None)
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
@ -767,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()`
@ -782,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:
@ -861,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()
@ -883,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.
@ -899,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

13
psycopg/psycopgmodule.c
View File

@ -257,7 +257,7 @@ psyco_connect(PyObject *self, PyObject *args, PyObject *keywds)
" * `conn_or_curs`: A connection, cursor or None"
#define typecast_from_python_doc \
"new_type(oids, name, adapter) -> new type object\n\n" \
"new_type(oids, name, castobj) -> new type object\n\n" \
"Create a new binding object. The object can be used with the\n" \
"`register_type()` function to bind PostgreSQL objects to python objects.\n\n" \
":Parameters:\n" \
@ -268,6 +268,15 @@ psyco_connect(PyObject *self, PyObject *args, PyObject *keywds)
" the string representation returned by PostgreSQL (`!None` if ``NULL``)\n" \
" and ``cur`` is the cursor from which data are read."
#define typecast_array_from_python_doc \
"new_array_type(oids, name, baseobj) -> new type object\n\n" \
"Create a new binding object to parse an array.\n\n" \
"The object can be used with `register_type()`.\n\n" \
":Parameters:\n" \
" * `oids`: Tuple of ``oid`` of the PostgreSQL types to convert.\n" \
" * `name`: Name for the new type\n" \
" * `baseobj`: Adapter to perform type conversion of a single array item."
static void
_psyco_register_type_set(PyObject **dict, PyObject *type)
{
@ -758,6 +767,8 @@ static PyMethodDef psycopgMethods[] = {
METH_VARARGS, psyco_register_type_doc},
{"new_type", (PyCFunction)typecast_from_python,
METH_VARARGS|METH_KEYWORDS, typecast_from_python_doc},
{"new_array_type", (PyCFunction)typecast_array_from_python,
METH_VARARGS|METH_KEYWORDS, typecast_array_from_python_doc},
{"AsIs", (PyCFunction)psyco_AsIs,
METH_VARARGS, psyco_AsIs_doc},

23
psycopg/typecast.c
View File

@ -603,6 +603,29 @@ typecast_from_python(PyObject *self, PyObject *args, PyObject *keywds)
return typecast_new(name, v, cast, base);
}
PyObject *
typecast_array_from_python(PyObject *self, PyObject *args, PyObject *keywds)
{
PyObject *values, *name = NULL, *base = NULL;
typecastObject *obj = NULL;
static char *kwlist[] = {"values", "name", "baseobj", NULL};
if (!PyArg_ParseTupleAndKeywords(args, keywds, "O!O!O!", kwlist,
&PyTuple_Type, &values,
&Text_Type, &name,
&typecastType, &base)) {
return NULL;
}
if ((obj = (typecastObject *)typecast_new(name, values, NULL, base))) {
obj->ccast = typecast_GENERIC_ARRAY_cast;
obj->pcast = NULL;
}
return (PyObject *)obj;
}
PyObject *
typecast_from_c(typecastObject_initlist *type, PyObject *dict)
{

4
psycopg/typecast.h
View File

@ -77,9 +77,11 @@ HIDDEN int typecast_add(PyObject *obj, PyObject *dict, int binary);
/* the C callable typecastObject creator function */
HIDDEN PyObject *typecast_from_c(typecastObject_initlist *type, PyObject *d);
/* the python callable typecast creator function */
/* the python callable typecast creator functions */
HIDDEN PyObject *typecast_from_python(
PyObject *self, PyObject *args, PyObject *keywds);
HIDDEN PyObject *typecast_array_from_python(
PyObject *self, PyObject *args, PyObject *keywds);
/* the function used to dispatch typecasting calls */
HIDDEN PyObject *typecast_cast(

9
psycopg/typecast_array.c
View File

@ -133,7 +133,7 @@ typecast_array_tokenize(const char *str, Py_ssize_t strlength,
}
if (res == ASCAN_QUOTED) {
Py_ssize_t j;
const char *j, *jj;
char *buffer = PyMem_Malloc(l+1);
if (buffer == NULL) {
PyErr_NoMemory();
@ -142,10 +142,9 @@ typecast_array_tokenize(const char *str, Py_ssize_t strlength,
*token = buffer;
for (j = *pos; j < *pos+l; j++) {
if (str[j] != '\\'
|| (j > *pos && str[j-1] == '\\'))
*(buffer++) = str[j];
for (j = str + *pos, jj = j + l; j < jj; ++j) {
if (*j == '\\') { ++j; }
*(buffer++) = *j;
}
*buffer = '\0';

24
tests/types_basic.py
View File

@ -189,6 +189,17 @@ class TypesBasicTests(unittest.TestCase):
s = self.execute("SELECT '{}'::text AS foo")
self.failUnlessEqual(s, "{}")
def testArrayEscape(self):
ss = ['', '\\', '"', '\\\\', '\\"']
for s in ss:
r = self.execute("SELECT %s AS foo", (s,))
self.failUnlessEqual(s, r)
r = self.execute("SELECT %s AS foo", ([s],))
self.failUnlessEqual([s], r)
r = self.execute("SELECT %s AS foo", (ss,))
self.failUnlessEqual(ss, r)
@testutils.skip_from_python(3)
def testTypeRoundtripBuffer(self):
o1 = buffer("".join(map(chr, range(256))))
@ -285,6 +296,19 @@ class TypesBasicTests(unittest.TestCase):
l1 = self.execute("select -%s;", (-1L,))
self.assertEqual(1, l1)
def testGenericArray(self):
def caster(s, cur):
if s is None: return "nada"
return int(s) * 2
base = psycopg2.extensions.new_type((23,), "INT4", caster)
array = psycopg2.extensions.new_array_type((1007,), "INT4ARRAY", base)
psycopg2.extensions.register_type(array, self.conn)
a = self.execute("select '{1,2,3}'::int4[]")
self.assertEqual(a, [2,4,6])
a = self.execute("select '{1,2,NULL}'::int4[]")
self.assertEqual(a, [2,4,'nada'])
class AdaptSubclassTest(unittest.TestCase):
def test_adapt_subtype(self):

82
tests/types_extras.py
View File

@ -22,7 +22,7 @@ import re
import sys
from datetime import date
from testutils import unittest, skip_if_no_uuid
from testutils import unittest, skip_if_no_uuid, skip_before_postgres
import psycopg2
import psycopg2.extras
@ -357,6 +357,63 @@ class HstoreTestCase(unittest.TestCase):
finally:
psycopg2.extensions.string_types.pop(oid)
@skip_if_no_hstore
@skip_before_postgres(8, 3)
def test_roundtrip_array(self):
from psycopg2.extras import register_hstore
register_hstore(self.conn)
ds = []
ds.append({})
ds.append({'a': 'b', 'c': None})
ab = map(chr, range(32, 128))
ds.append(dict(zip(ab, ab)))
ds.append({''.join(ab): ''.join(ab)})
self.conn.set_client_encoding('latin1')
if sys.version_info[0] < 3:
ab = map(chr, range(32, 127) + range(160, 255))
else:
ab = bytes(range(32, 127) + range(160, 255)).decode('latin1')
ds.append({''.join(ab): ''.join(ab)})
ds.append(dict(zip(ab, ab)))
cur = self.conn.cursor()
cur.execute("select %s", (ds,))
ds1 = cur.fetchone()[0]
self.assertEqual(ds, ds1)
@skip_if_no_hstore
@skip_before_postgres(8, 3)
def test_array_cast(self):
from psycopg2.extras import register_hstore
register_hstore(self.conn)
cur = self.conn.cursor()
cur.execute("select array['a=>1'::hstore, 'b=>2'::hstore];")
a = cur.fetchone()[0]
self.assertEqual(a, [{'a': '1'}, {'b': '2'}])
@skip_if_no_hstore
def test_array_cast_oid(self):
cur = self.conn.cursor()
cur.execute("select 'hstore'::regtype::oid, 'hstore[]'::regtype::oid")
oid, aoid = cur.fetchone()
from psycopg2.extras import register_hstore
register_hstore(None, globally=True, oid=oid, array_oid=aoid)
try:
cur.execute("select null::hstore, ''::hstore, 'a => b'::hstore, '{a=>b}'::hstore[]")
t = cur.fetchone()
self.assert_(t[0] is None)
self.assertEqual(t[1], {})
self.assertEqual(t[2], {'a': 'b'})
self.assertEqual(t[3], [{'a': 'b'}])
finally:
psycopg2.extensions.string_types.pop(oid)
psycopg2.extensions.string_types.pop(aoid)
def skip_if_no_composite(f):
def skip_if_no_composite_(self):
@ -563,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: