From 6bca443e37cd0cb5e37352eaf2e8aafc832324ee Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Wed, 13 Aug 2014 00:43:33 +0100
Subject: [PATCH] Added name param to register_json()

---
 doc/src/extras.rst |  3 +++
 lib/_json.py       | 28 +++++++++++++++-------------
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/doc/src/extras.rst b/doc/src/extras.rst
index 7fab3384..0abd3354 100644
--- a/doc/src/extras.rst
+++ b/doc/src/extras.rst
@@ -248,6 +248,9 @@ from :sql:`json` into :py:class:`~decimal.Decimal` you can use::
 
 .. autofunction:: register_json
 
+    .. versionchanged:: 2.5.4
+        added the *name* parameter to enable :sql:`jsonb` support.
+
 .. autofunction:: register_default_json
 
 
diff --git a/lib/_json.py b/lib/_json.py
index 3a4361e8..047551e5 100644
--- a/lib/_json.py
+++ b/lib/_json.py
@@ -102,7 +102,7 @@ class Json(object):
 
 
 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.
 
     :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*
     :param array_oid: the OID of the :sql:`json[]` array type if known;
         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
-    database and look for the OID of the :sql:`json` type. No query is
-    performed if *oid* and *array_oid* are provided.  Raise
-    `~psycopg2.ProgrammingError` if the type is not found.
+    database and look for the OID of the :sql:`json` type (or an alternative
+    type if *name* if provided). No query is performed if *oid* and *array_oid*
+    are provided.  Raise `~psycopg2.ProgrammingError` if the type is not found.
 
     """
     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)
 
@@ -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,
         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."""
     if loads is None:
         if json is None:
@@ -162,15 +164,15 @@ def _create_json_typecasters(oid, array_oid, loads=None):
             return None
         return loads(s)
 
-    JSON = new_type((oid, ), 'JSON', typecast_json)
+    JSON = new_type((oid, ), name, typecast_json)
     if array_oid is not None:
-        JSONARRAY = new_array_type((array_oid, ), "JSONARRAY", JSON)
+        JSONARRAY = new_array_type((array_oid, ), "%sARRAY" % name, JSON)
     else:
         JSONARRAY = None
 
     return JSON, JSONARRAY
 
-def _get_json_oids(conn_or_curs):
+def _get_json_oids(conn_or_curs, name='json'):
     # lazy imports
     from psycopg2.extensions import STATUS_IN_TRANSACTION
     from psycopg2.extras import _solve_conn_curs
@@ -185,8 +187,8 @@ def _get_json_oids(conn_or_curs):
 
     # get the oid for the hstore
     curs.execute(
-        "SELECT t.oid, %s FROM pg_type t WHERE t.typname = 'json';"
-            % typarray)
+        "SELECT t.oid, %s FROM pg_type t WHERE t.typname = %%s;"
+            % typarray, (name,))
     r = curs.fetchone()
 
     # revert the status of the connection as before the command
@@ -194,7 +196,7 @@ def _get_json_oids(conn_or_curs):
         conn.rollback()
 
     if not r:
-        raise conn.ProgrammingError("json data type not found")
+        raise conn.ProgrammingError("%s data type not found" % name)
 
     return r