- merge of ticket_3499 indexed access branch

- The "hashable" flag on special datatypes such as :class:`.postgresql.ARRAY`,
:class:`.postgresql.JSON` and :class:`.postgresql.HSTORE` is now
set to False, which allows these types to be fetchable in ORM
queries that include entities within the row.  fixes #3499
- The Postgresql :class:`.postgresql.ARRAY` type now supports multidimensional
indexed access, e.g. expressions such as ``somecol[5][6]`` without
any need for explicit casts or type coercions, provided
that the :paramref:`.postgresql.ARRAY.dimensions` parameter is set to the
desired number of dimensions. fixes #3487
- The return type for the :class:`.postgresql.JSON` and :class:`.postgresql.JSONB`
when using indexed access has been fixed to work like Postgresql itself,
and returns an expression that itself is of type :class:`.postgresql.JSON`
or :class:`.postgresql.JSONB`.  Previously, the accessor would return
:class:`.NullType` which disallowed subsequent JSON-like operators to be
used. part of fixes #3503
- The :class:`.postgresql.JSON`, :class:`.postgresql.JSONB` and
:class:`.postgresql.HSTORE` datatypes now allow full control over the
return type from an indexed textual access operation, either ``column[someindex].astext``
for a JSON type or ``column[someindex]`` for an HSTORE type,
via the :paramref:`.postgresql.JSON.astext_type` and
:paramref:`.postgresql.HSTORE.text_type` parameters. also part of fixes #3503
- The :attr:`.postgresql.JSON.Comparator.astext` modifier no longer
calls upon :meth:`.ColumnElement.cast` implicitly, as PG's JSON/JSONB
types allow cross-casting between each other as well.  Code that
makes use of :meth:`.ColumnElement.cast` on JSON indexed access,
e.g. ``col[someindex].cast(Integer)``, will need to be changed
to call :attr:`.postgresql.JSON.Comparator.astext` explicitly.  This is
part of the refactor in references #3503 for consistency in operator
use.

1 files changed, 137 insertions, 166 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/json.py b/lib/sqlalchemy/dialects/postgresql/json.py
index 13ebc4afe..4716ca970 100644
--- a/lib/sqlalchemy/dialects/postgresql/json.py
+++ b/lib/sqlalchemy/dialects/postgresql/json.py
@@ -6,96 +6,60 @@
 # the MIT License: http://www.opensource.org/licenses/mit-license.php
 from __future__ import absolute_import
 
+import collections
 import json
 
 from .base import ischema_names
 from ... import types as sqltypes
-from ...sql.operators import custom_op
-from ... import sql
-from ...sql import elements, default_comparator
+from ...sql import operators
+from ...sql import elements
 from ... import util
 
-__all__ = ('JSON', 'JSONElement', 'JSONB')
+__all__ = ('JSON', 'JSONB')
 
 
-class JSONElement(elements.BinaryExpression):
-    """Represents accessing an element of a :class:`.JSON` value.
+# json : returns json
+INDEX = operators.custom_op(
+    "->", precedence=5, natural_self_precedent=True
+)
 
-    The :class:`.JSONElement` is produced whenever using the Python index
-    operator on an expression that has the type :class:`.JSON`::
+# path operator: returns json
+PATHIDX = operators.custom_op(
+    "#>", precedence=5, natural_self_precedent=True
+)
 
-        expr = mytable.c.json_data['some_key']
+# json + astext: returns text
+ASTEXT = operators.custom_op(
+    "->>", precedence=5, natural_self_precedent=True
+)
 
-    The expression typically compiles to a JSON access such as ``col -> key``.
-    Modifiers are then available for typing behavior, including
-    :meth:`.JSONElement.cast` and :attr:`.JSONElement.astext`.
+# path operator  + astext: returns text
+ASTEXT_PATHIDX = operators.custom_op(
+    "#>>", precedence=5, natural_self_precedent=True
+)
 
-    """
-
-    def __init__(self, left, right, astext=False,
-                 opstring=None, result_type=None):
-        self._astext = astext
-        if opstring is None:
-            if hasattr(right, '__iter__') and \
-                    not isinstance(right, util.string_types):
-                opstring = "#>"
-                right = "{%s}" % (
-                    ", ".join(util.text_type(elem) for elem in right))
-            else:
-                opstring = "->"
-
-        self._json_opstring = opstring
-        operator = custom_op(opstring, precedence=5)
-        right = default_comparator._check_literal(
-            left, operator, right)
-        super(JSONElement, self).__init__(
-            left, right, operator, type_=result_type)
-
-    @property
-    def astext(self):
-        """Convert this :class:`.JSONElement` to use the 'astext' operator
-        when evaluated.
-
-        E.g.::
-
-            select([data_table.c.data['some key'].astext])
-
-        .. seealso::
-
-            :meth:`.JSONElement.cast`
+HAS_KEY = operators.custom_op(
+    "?", precedence=5, natural_self_precedent=True
+)
 
-        """
-        if self._astext:
-            return self
-        else:
-            return JSONElement(
-                self.left,
-                self.right,
-                astext=True,
-                opstring=self._json_opstring + ">",
-                result_type=sqltypes.String(convert_unicode=True)
-            )
-
-    def cast(self, type_):
-        """Convert this :class:`.JSONElement` to apply both the 'astext' operator
-        as well as an explicit type cast when evaluated.
+HAS_ALL = operators.custom_op(
+    "?&", precedence=5, natural_self_precedent=True
+)
 
-        E.g.::
+HAS_ANY = operators.custom_op(
+    "?|", precedence=5, natural_self_precedent=True
+)
 
-            select([data_table.c.data['some key'].cast(Integer)])
+CONTAINS = operators.custom_op(
+    "@>", precedence=5, natural_self_precedent=True
+)
 
-        .. seealso::
+CONTAINED_BY = operators.custom_op(
+    "<@", precedence=5, natural_self_precedent=True
+)
 
-            :attr:`.JSONElement.astext`
-
-        """
-        if not self._astext:
-            return self.astext.cast(type_)
-        else:
-            return sql.cast(self, type_)
 
-
-class JSON(sqltypes.TypeEngine):
+class JSON(sqltypes.Indexable, sqltypes.TypeEngine):
     """Represent the Postgresql JSON type.
 
     The :class:`.JSON` type stores arbitrary JSON format data, e.g.::
@@ -113,31 +77,36 @@ class JSON(sqltypes.TypeEngine):
 
     :class:`.JSON` provides several operations:
 
-    * Index operations::
+    * Index operations (the ``->`` operator)::
 
         data_table.c.data['some key']
 
-    * Index operations returning text (required for text comparison)::
+    * Index operations returning text (the ``->>`` operator)::
 
         data_table.c.data['some key'].astext == 'some value'
 
-    * Index operations with a built-in CAST call::
+    * Index operations with CAST
+      (equivalent to ``CAST(col ->> ['some key'] AS <type>)``)::
 
-        data_table.c.data['some key'].cast(Integer) == 5
+        data_table.c.data['some key'].astext.cast(Integer) == 5
 
-    * Path index operations::
+    * Path index operations (the ``#>`` operator)::
 
         data_table.c.data[('key_1', 'key_2', ..., 'key_n')]
 
-    * Path index operations returning text (required for text comparison)::
+    * Path index operations returning text (the ``#>>`` operator)::
+
+        data_table.c.data[('key_1', 'key_2', ..., 'key_n')].astext == \
+'some value'
 
-        data_table.c.data[('key_1', 'key_2', ..., 'key_n')].astext == \\
-            'some value'
+    .. versionchanged:: 1.1  The :meth:`.ColumnElement.cast` operator on
+       JSON objects now requires that the :attr:`.JSON.Comparator.astext`
+       modifier be called explicitly, if the cast works only from a textual
+       string.
 
-    Index operations return an instance of :class:`.JSONElement`, which
-    represents an expression such as ``column -> index``.  This element then
-    defines methods such as :attr:`.JSONElement.astext` and
-    :meth:`.JSONElement.cast` for setting up type behavior.
+    Index operations return an expression object whose type defaults to
+    :class:`.JSON` by default, so that further JSON-oriented instructions
+    may be called upon the result type.
 
     The :class:`.JSON` type, when used with the SQLAlchemy ORM, does not
     detect in-place mutations to the structure.  In order to detect these, the
@@ -161,11 +130,18 @@ class JSON(sqltypes.TypeEngine):
 
     .. versionadded:: 0.9
 
+    .. seealso::
+
+        :class:`.JSONB`
+
     """
 
     __visit_name__ = 'JSON'
 
-    def __init__(self, none_as_null=False):
+    hashable = False
+    astext_type = sqltypes.Text()
+
+    def __init__(self, none_as_null=False, astext_type=None):
         """Construct a :class:`.JSON` type.
 
         :param none_as_null: if True, persist the value ``None`` as a
@@ -179,23 +155,56 @@ class JSON(sqltypes.TypeEngine):
          .. versionchanged:: 0.9.8 - Added ``none_as_null``, and :func:`.null`
             is now supported in order to persist a NULL value.
 
+        :param astext_type: the type to use for the
+         :attr:`.JSON.Comparator.astext`
+         accessor on indexed attributes.  Defaults to :class:`.types.Text`.
+
+         .. versionadded:: 1.1.0
+
          """
         self.none_as_null = none_as_null
+        if astext_type is not None:
+            self.astext_type = astext_type
 
-    class comparator_factory(sqltypes.Concatenable.Comparator):
+    class Comparator(
+            sqltypes.Indexable.Comparator, sqltypes.Concatenable.Comparator):
         """Define comparison operations for :class:`.JSON`."""
 
-        def __getitem__(self, other):
-            """Get the value at a given key."""
+        @property
+        def astext(self):
+            """On an indexed expression, use the "astext" (e.g. "->>")
+            conversion when rendered in SQL.
+
+            E.g.::
+
+                select([data_table.c.data['some key'].astext])
+
+            .. seealso::
+
+                :meth:`.ColumnElement.cast`
+
+            """
+            against = self.expr.operator
+            if against is PATHIDX:
+                against = ASTEXT_PATHIDX
+            else:
+                against = ASTEXT
+
+            return self.expr.left.operate(
+                against, self.expr.right, result_type=self.type.astext_type)
+
+        def _setup_getitem(self, index):
+            if not isinstance(index, util.string_types):
+                assert isinstance(index, collections.Sequence)
+                tokens = [util.text_type(elem) for elem in index]
+                index = "{%s}" % (", ".join(tokens))
+                operator = PATHIDX
+            else:
+                operator = INDEX
 
-            return JSONElement(self.expr, other)
+            return operator, index, self.type
 
-        def _adapt_expression(self, op, other_comparator):
-            if isinstance(op, custom_op):
-                if op.opstring == '->':
-                    return op, sqltypes.Text
-            return sqltypes.Concatenable.Comparator.\
-                _adapt_expression(self, op, other_comparator)
+    comparator_factory = Comparator
 
     def bind_processor(self, dialect):
         json_serializer = dialect._json_serializer or json.dumps
@@ -253,106 +262,68 @@ class JSONB(JSON):
                 data = {"key1": "value1", "key2": "value2"}
             )
 
-    :class:`.JSONB` provides several operations:
-
-    * Index operations::
-
-        data_table.c.data['some key']
-
-    * Index operations returning text (required for text comparison)::
+    The :class:`.JSONB` type includes all operations provided by
+    :class:`.JSON`, including the same behaviors for indexing operations.
+    It also adds additional operators specific to JSONB, including
+    :meth:`.JSONB.Comparator.has_key`, :meth:`.JSONB.Comparator.has_all`,
+    :meth:`.JSONB.Comparator.has_any`, :meth:`.JSONB.Comparator.contains`,
+    and :meth:`.JSONB.Comparator.contained_by`.
+
+    Like the :class:`.JSON` type, the :class:`.JSONB` type does not detect
+    in-place changes when used with the ORM, unless the
+    :mod:`sqlalchemy.ext.mutable` extension is used.
+
+    Custom serializers and deserializers
+    are shared with the :class:`.JSON` class, using the ``json_serializer``
+    and ``json_deserializer`` keyword arguments.  These must be specified
+    at the dialect level using :func:`.create_engine`.  When using
+    psycopg2, the serializers are associated with the jsonb type using
+    ``psycopg2.extras.register_default_jsonb`` on a per-connection basis,
+    in the same way that ``psycopg2.extras.register_default_json`` is used
+    to register these handlers with the json type.
 
-        data_table.c.data['some key'].astext == 'some value'
-
-    * Index operations with a built-in CAST call::
-
-        data_table.c.data['some key'].cast(Integer) == 5
-
-    * Path index operations::
-
-        data_table.c.data[('key_1', 'key_2', ..., 'key_n')]
-
-    * Path index operations returning text (required for text comparison)::
-
-        data_table.c.data[('key_1', 'key_2', ..., 'key_n')].astext == \\
-            'some value'
-
-    Index operations return an instance of :class:`.JSONElement`, which
-    represents an expression such as ``column -> index``.  This element then
-    defines methods such as :attr:`.JSONElement.astext` and
-    :meth:`.JSONElement.cast` for setting up type behavior.
-
-    The :class:`.JSON` type, when used with the SQLAlchemy ORM, does not
-    detect in-place mutations to the structure.  In order to detect these, the
-    :mod:`sqlalchemy.ext.mutable` extension must be used.  This extension will
-    allow "in-place" changes to the datastructure to produce events which
-    will be detected by the unit of work.  See the example at :class:`.HSTORE`
-    for a simple example involving a dictionary.
-
-    Custom serializers and deserializers are specified at the dialect level,
-    that is using :func:`.create_engine`.  The reason for this is that when
-    using psycopg2, the DBAPI only allows serializers at the per-cursor
-    or per-connection level.   E.g.::
+    .. versionadded:: 0.9.7
 
-        engine = create_engine("postgresql://scott:tiger@localhost/test",
-                                json_serializer=my_serialize_fn,
-                                json_deserializer=my_deserialize_fn
-                        )
+    .. seealso::
 
-    When using the psycopg2 dialect, the json_deserializer is registered
-    against the database using ``psycopg2.extras.register_default_json``.
-
-    .. versionadded:: 0.9.7
+        :class:`.JSON`
 
     """
 
     __visit_name__ = 'JSONB'
-    hashable = False
 
-    class comparator_factory(sqltypes.Concatenable.Comparator):
+    class Comparator(JSON.Comparator):
         """Define comparison operations for :class:`.JSON`."""
 
-        def __getitem__(self, other):
-            """Get the value at a given key."""
-
-            return JSONElement(self.expr, other)
-
-        def _adapt_expression(self, op, other_comparator):
-            # How does one do equality?? jsonb also has "=" eg.
-            # '[1,2,3]'::jsonb = '[1,2,3]'::jsonb
-            if isinstance(op, custom_op):
-                if op.opstring in ['?', '?&', '?|', '@>', '<@']:
-                    return op, sqltypes.Boolean
-                if op.opstring == '->':
-                    return op, sqltypes.Text
-            return sqltypes.Concatenable.Comparator.\
-                _adapt_expression(self, op, other_comparator)
-
         def has_key(self, other):
             """Boolean expression.  Test for presence of a key.  Note that the
             key may be a SQLA expression.
             """
-            return self.expr.op('?')(other)
+            return self.operate(HAS_KEY, other, result_type=sqltypes.Boolean)
 
         def has_all(self, other):
             """Boolean expression.  Test for presence of all keys in jsonb
             """
-            return self.expr.op('?&')(other)
+            return self.operate(HAS_ALL, other, result_type=sqltypes.Boolean)
 
         def has_any(self, other):
             """Boolean expression.  Test for presence of any key in jsonb
             """
-            return self.expr.op('?|')(other)
+            return self.operate(HAS_ANY, other, result_type=sqltypes.Boolean)
 
         def contains(self, other, **kwargs):
-            """Boolean expression.  Test if keys (or array) are a superset of/contained
-            the keys of the argument jsonb expression.
+            """Boolean expression.  Test if keys (or array) are a superset
+            of/contained the keys of the argument jsonb expression.
             """
-            return self.expr.op('@>')(other)
+            return self.operate(CONTAINS, other, result_type=sqltypes.Boolean)
 
         def contained_by(self, other):
             """Boolean expression.  Test if keys are a proper subset of the
             keys of the argument jsonb expression.
             """
-            return self.expr.op('<@')(other)
+            return self.operate(
+                CONTAINED_BY, other, result_type=sqltypes.Boolean)
+
+    comparator_factory = Comparator
 
 ischema_names['jsonb'] = JSONB