Added name param to register_json()

This commit is contained in:
Daniele Varrazzo 2014-08-13 00:43:33 +01:00
parent 634fc004fb
commit 6bca443e37
2 changed files with 18 additions and 13 deletions

View File

@ -248,6 +248,9 @@ from :sql:`json` into :py:class:`~decimal.Decimal` you can use::
.. autofunction:: register_json .. autofunction:: register_json
.. versionchanged:: 2.5.4
added the *name* parameter to enable :sql:`jsonb` support.
.. autofunction:: register_default_json .. autofunction:: register_default_json

View File

@ -102,7 +102,7 @@ class Json(object):
def register_json(conn_or_curs=None, globally=False, loads=None, def register_json(conn_or_curs=None, globally=False, loads=None,
oid=None, array_oid=None): oid=None, array_oid=None, name='json'):
"""Create and register typecasters converting :sql:`json` type to Python objects. """Create and register typecasters converting :sql:`json` type to Python objects.
:param conn_or_curs: a connection or cursor used to find the :sql:`json` :param conn_or_curs: a connection or cursor used to find the :sql:`json`
@ -118,17 +118,19 @@ def register_json(conn_or_curs=None, globally=False, loads=None,
queried on *conn_or_curs* queried on *conn_or_curs*
:param array_oid: the OID of the :sql:`json[]` array type if known; :param array_oid: the OID of the :sql:`json[]` array type if known;
if not, it will be queried on *conn_or_curs* if not, it will be queried on *conn_or_curs*
:param name: the name of the data type to look for in *conn_or_curs*
The connection or cursor passed to the function will be used to query the The connection or cursor passed to the function will be used to query the
database and look for the OID of the :sql:`json` type. No query is database and look for the OID of the :sql:`json` type (or an alternative
performed if *oid* and *array_oid* are provided. Raise type if *name* if provided). No query is performed if *oid* and *array_oid*
`~psycopg2.ProgrammingError` if the type is not found. are provided. Raise `~psycopg2.ProgrammingError` if the type is not found.
""" """
if oid is None: if oid is None:
oid, array_oid = _get_json_oids(conn_or_curs) oid, array_oid = _get_json_oids(conn_or_curs, name)
JSON, JSONARRAY = _create_json_typecasters(oid, array_oid, loads) JSON, JSONARRAY = _create_json_typecasters(
oid, array_oid, loads=loads, name=name.upper())
register_type(JSON, not globally and conn_or_curs or None) register_type(JSON, not globally and conn_or_curs or None)
@ -149,7 +151,7 @@ def register_default_json(conn_or_curs=None, globally=False, loads=None):
return register_json(conn_or_curs=conn_or_curs, globally=globally, return register_json(conn_or_curs=conn_or_curs, globally=globally,
loads=loads, oid=JSON_OID, array_oid=JSONARRAY_OID) loads=loads, oid=JSON_OID, array_oid=JSONARRAY_OID)
def _create_json_typecasters(oid, array_oid, loads=None): def _create_json_typecasters(oid, array_oid, loads=None, name='JSON'):
"""Create typecasters for json data type.""" """Create typecasters for json data type."""
if loads is None: if loads is None:
if json is None: if json is None:
@ -162,15 +164,15 @@ def _create_json_typecasters(oid, array_oid, loads=None):
return None return None
return loads(s) return loads(s)
JSON = new_type((oid, ), 'JSON', typecast_json) JSON = new_type((oid, ), name, typecast_json)
if array_oid is not None: if array_oid is not None:
JSONARRAY = new_array_type((array_oid, ), "JSONARRAY", JSON) JSONARRAY = new_array_type((array_oid, ), "%sARRAY" % name, JSON)
else: else:
JSONARRAY = None JSONARRAY = None
return JSON, JSONARRAY return JSON, JSONARRAY
def _get_json_oids(conn_or_curs): def _get_json_oids(conn_or_curs, name='json'):
# lazy imports # lazy imports
from psycopg2.extensions import STATUS_IN_TRANSACTION from psycopg2.extensions import STATUS_IN_TRANSACTION
from psycopg2.extras import _solve_conn_curs from psycopg2.extras import _solve_conn_curs
@ -185,8 +187,8 @@ def _get_json_oids(conn_or_curs):
# get the oid for the hstore # get the oid for the hstore
curs.execute( curs.execute(
"SELECT t.oid, %s FROM pg_type t WHERE t.typname = 'json';" "SELECT t.oid, %s FROM pg_type t WHERE t.typname = %%s;"
% typarray) % typarray, (name,))
r = curs.fetchone() r = curs.fetchone()
# revert the status of the connection as before the command # revert the status of the connection as before the command
@ -194,7 +196,7 @@ def _get_json_oids(conn_or_curs):
conn.rollback() conn.rollback()
if not r: if not r:
raise conn.ProgrammingError("json data type not found") raise conn.ProgrammingError("%s data type not found" % name)
return r return r