- most of the index model working, no tests yet

- ARRAY restored, most of JSON, not JSONB yet

4 files changed, 252 insertions, 65 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/array.py b/lib/sqlalchemy/dialects/postgresql/array.py
index c5bcdebd0..2e27d9bd6 100644
--- a/lib/sqlalchemy/dialects/postgresql/array.py
+++ b/lib/sqlalchemy/dialects/postgresql/array.py
@@ -169,15 +169,25 @@ class ARRAY(sqltypes.Indexable, sqltypes.Concatenable, sqltypes.TypeEngine):
             mytable.c.data[2:7]: [1, 2, 3]
         })
 
-    .. note::
+    Multi-dimensional array index support is provided automatically based on
+    either the value specified for the :paramref:`.ARRAY.dimensions` parameter,
+    or more specifically the :paramref:`.ARRAY.index_map` parameter.
+    E.g. an :class:`.ARRAY` with dimensions set to 2 would return an expression
+    of type :class:`.ARRAY` for a single index operation::
 
-        Multi-dimensional support for the ``[]`` operator is not supported
-        in SQLAlchemy 1.0.  Please use the :func:`.type_coerce` function
-        to cast an intermediary expression to ARRAY again as a workaround::
+        type = ARRAY(Integer, dimensions=2)
 
-            expr = type_coerce(my_array_column[5], ARRAY(Integer))[6]
+        expr = column('x', type)  # expr is of type ARRAY(Integer, dimensions=2)
 
-        Multi-dimensional support will be provided in a future release.
+        expr = column('x', type)[5]  # expr is of type ARRAY(Integer, dimensions=1)
+
+    An index expression from ``expr`` above would then return an expression
+    of type Integer::
+
+        sub_expr = expr[10]  # expr is of type Integer
+
+    .. versionadded:: 1.1 support for index operations on multi-dimensional
+       :class:`.postgresql.ARRAY` objects is added.
 
     :class:`.ARRAY` provides special methods for containment operations,
     e.g.::
@@ -205,6 +215,13 @@ class ARRAY(sqltypes.Indexable, sqltypes.Concatenable, sqltypes.TypeEngine):
 
         """Define comparison operations for :class:`.ARRAY`."""
 
+        def _type_for_index(self, index):
+            new_type = super(ARRAY.Comparator, self)._type_for_index(index)
+            if self.type.dimensions is not None:
+                new_type.dimensions = self.dimensions - 1
+
+            return new_type
+
         def any(self, other, operator=operators.eq):
             """Return ``other operator ANY (array)`` clause.
 
@@ -269,34 +286,25 @@ class ARRAY(sqltypes.Indexable, sqltypes.Concatenable, sqltypes.TypeEngine):
             """Boolean expression.  Test if elements are a superset of the
             elements of the argument array expression.
             """
-            return self.operate(CONTAINS, other)
+            return self.operate(CONTAINS, other, result_type=sqltypes.Boolean)
 
         def contained_by(self, other):
             """Boolean expression.  Test if elements are a proper subset of the
             elements of the argument array expression.
             """
-            return self.operate(CONTAINED_BY, other)
+            return self.operate(
+                CONTAINED_BY, other, result_type=sqltypes.Boolean)
 
         def overlap(self, other):
             """Boolean expression.  Test if array has elements in common with
             an argument array expression.
             """
-            return self.operate(OVERLAP, other)
-
-        def _index_map_type(self, right_comparator):
-            return self.type
-
-        def _adapt_expression(self, op, other_comparator):
-            if op in (CONTAINS, CONTAINED_BY, OVERLAP):
-                return op, sqltypes.Boolean
-            else:
-                return super(ARRAY.Comparator, self).\
-                    _adapt_expression(op, other_comparator)
+            return self.operate(OVERLAP, other, result_type=sqltypes.Boolean)
 
     comparator_factory = Comparator
 
     def __init__(self, item_type, as_tuple=False, dimensions=None,
-                 zero_indexes=False):
+                 zero_indexes=False, index_map=None):
         """Construct an ARRAY.
 
         E.g.::
@@ -330,6 +338,26 @@ class ARRAY(sqltypes.Indexable, sqltypes.Concatenable, sqltypes.TypeEngine):
 
          .. versionadded:: 0.9.5
 
+        :param index_map: type map used by the getitem operator, e.g.
+         expressions like ``col[5]``.  See :class:`.Indexable` for a
+         description of how this map is configured.
+
+         For the :class:`.ARRAY` class, this map if omitted is
+         automatically configured based on the  number of dimensions
+         given, that is an :class:`.ARRAY` that specifies
+         ``dimensions=3`` and a return type of ``String`` would
+         generate an  index_map of  ``{ANY_KEY: {ANY_KEY: {ANY_KEY:
+         String}}}``, so that an expression of ``col[x][y][z]`` would
+         yield arrays until the last index, which  yields a string
+         expression.
+
+         When the map and the dimensions argument is omitted, the map here
+         assumes a 1-dimensional array and defaults to
+         ``{ANY_KEY: <array type}``.
+
+         ..versionadded:: 1.1
+
+
         """
         if isinstance(item_type, ARRAY):
             raise ValueError("Do not nest ARRAY types; ARRAY(basetype) "
@@ -340,6 +368,20 @@ class ARRAY(sqltypes.Indexable, sqltypes.Concatenable, sqltypes.TypeEngine):
         self.as_tuple = as_tuple
         self.dimensions = dimensions
         self.zero_indexes = zero_indexes
+        if index_map is not None:
+            self.index_map = index_map
+        elif self.dimensions is not None:
+            self.index_map = d = {
+                ARRAY.SLICE_TYPE: ARRAY.SAME_TYPE
+            }
+            for i in range(self.dimensions):
+                d[ARRAY.ANY_KEY] = d = {}
+            d[ARRAY.ANY_KEY] = self.item_type
+        else:
+            self.index_map = {
+                ARRAY.ANY_KEY: self.item_type,
+                ARRAY.SLICE_TYPE: ARRAY.SAME_TYPE
+            }
 
     @property
     def hashable(self):
@@ -405,4 +447,4 @@ class ARRAY(sqltypes.Indexable, sqltypes.Concatenable, sqltypes.TypeEngine):
                     tuple if self.as_tuple else list)
         return process
 
-ischema_names['_array'] = ARRAY
\ No newline at end of file
+ischema_names['_array'] = ARRAY
diff --git a/lib/sqlalchemy/dialects/postgresql/json.py b/lib/sqlalchemy/dialects/postgresql/json.py
index 8b1a70820..04480845c 100644
--- a/lib/sqlalchemy/dialects/postgresql/json.py
+++ b/lib/sqlalchemy/dialects/postgresql/json.py
@@ -111,7 +111,7 @@ class JSON(sqltypes.Indexable, sqltypes.TypeEngine):
 
     hashable = False
 
-    def __init__(self, none_as_null=False):
+    def __init__(self, none_as_null=False, index_map=None):
         """Construct a :class:`.JSON` type.
 
         :param none_as_null: if True, persist the value ``None`` as a
@@ -125,23 +125,23 @@ class JSON(sqltypes.Indexable, sqltypes.TypeEngine):
          .. versionchanged:: 0.9.8 - Added ``none_as_null``, and :func:`.null`
             is now supported in order to persist a NULL value.
 
+        :param index_map: type map used by the getitem operator, e.g.
+         expressions like ``col[5]``.  See :class:`.Indexable` for a
+         description of how this map is configured.   The index_map
+         for the :class:`.JSON` and :class:`.JSONB` types defaults to
+         ``{ANY_KEY: SAME_TYPE}``.
+
+         .. versionadded: 1.1
+
          """
         self.none_as_null = none_as_null
+        if index_map is not None:
+            self.index_map = index_map
 
-    class Comparator(sqltypes.Concatenable.Comparator):
+    class Comparator(
+            sqltypes.Indexable.Comparator, sqltypes.Concatenable.Comparator):
         """Define comparison operations for :class:`.JSON`."""
 
-        def __init__(self, expr, astext=False, aspath=False):
-            super(JSON.comparator_factory, self).__init__(expr)
-            self._astext = astext
-            self._aspath = aspath
-
-        def _clone(self, astext=False, aspath=False):
-            return self.__class__(
-                self.expr,
-                astext=self._astext or astext,
-                aspath=self._aspath or aspath)
-
         @property
         def astext(self):
             """On an indexed expression, use the "astext" (e.g. "->>")
@@ -156,37 +156,25 @@ class JSON(sqltypes.Indexable, sqltypes.TypeEngine):
                 :meth:`.ColumnElement.cast`
 
             """
-            if self._astext:
-                return self
+            against = self.expr.operator
+            if against is PATHIDX:
+                against = ASTEXT_PATHIDX
             else:
-                return self._clone(astext=True)
+                against = ASTEXT
 
-        def __getitem__(self, index):
-            if isinstance(index, collections.Sequence):
+            return self.expr.left.operate(
+                against, self.expr.right, result_type=sqltypes.Text)
+
+        def _setup_getitem(self, index):
+            if not isinstance(index, util.string_types):
+                assert isinstance(index, collections.Sequence)
                 index = "{%s}" % (
                     ", ".join(util.text_type(elem) for elem in index))
-                aspath = True
+                operator = PATHIDX
             else:
-                aspath = False
-            return self.operate(operators.getitem, index, aspath=aspath)
+                operator = INDEX
 
-        def _adapt_expression(self, op, other_comparator):
-            if op is operators.getitem:
-                if self._astext:
-                    if self._aspath:
-                        return ASTEXT_PATHIDX, sqltypes.Text
-                    else:
-                        return ASTEXT, sqltypes.Text
-                else:
-                    if self._aspath:
-                        # TODO: consult index map
-                        return PATHIDX, self.type
-                    else:
-                        # TODO: consult index map
-                        return INDEX, self.type
-            else:
-                return super(JSON.comparator_factory, self)._adapt_expression(
-                    op, other_comparator)
+            return operator, index, self._type_for_index(index)
 
     comparator_factory = Comparator
 
diff --git a/lib/sqlalchemy/sql/sqltypes.py b/lib/sqlalchemy/sql/sqltypes.py
index 1f6228959..e4b73e14f 100644
--- a/lib/sqlalchemy/sql/sqltypes.py
+++ b/lib/sqlalchemy/sql/sqltypes.py
@@ -9,6 +9,7 @@
 
 """
 
+import collections
 import datetime as dt
 import codecs
 
@@ -68,7 +69,7 @@ class Concatenable(object):
                     )):
                 return operators.concat_op, self.expr.type
             else:
-                return super(Concatenable.Comparator)._adapt_expression(
+                return super(Concatenable.Comparator, self)._adapt_expression(
                     op, other_comparator)
 
     comparator_factory = Comparator
@@ -78,22 +79,176 @@ class Indexable(object):
     """A mixin that marks a type as supporting indexing operations,
     such as array or JSON structures.
 
+    The key feature provided by :class:`.Indexable` is the index map feature.
+    This allows SQL expressions that use the index operator to have
+    an expected return type based on the structure.   For example, if
+    a JSON type expects to return json structures with the keys "a" and "b",
+    and the value referred to by "a" is an integer, and the value referred
+    to by "b" is itself a JSON structure with keys "c" and "d" referring
+    to String values, an index map for this schema would look like the
+    following::
+
+        # assume JSON is an Indexable subclass
+
+        type = JSON(index_map={
+                "a": Integer,
+                "b": {
+                    "c": String,
+                    "d": String
+                }
+            }
+        )
+
+    An expression that uses this type would indicate the following return
+    types::
+
+        column('x', type)['a']  # returns Integer
+        column('x', type)['b']  # returns JSON(index_map={"c": String, "d": String})
+        column('x', type)['b']['c']  # returns String
+        column('x', type)['unknown'] # raises an error
+
+    To indicate any key returned, the symbol :attr:`.Indexable.ANY_KEY` may
+    be used::
+
+        type = JSON(index_map={
+                JSON.ANY_KEY: {
+                    "x": Integer,
+                    "y": Integer
+                }
+            }
+        )
+
+    The above structure would be appropriate for a list of two-element mappings,
+    where any value in the first position, including SQL expressions, returns
+    the structure::
+
+        column('x', type)[5]['x']  # returns Integer
+        column('x', type)[column('y')]['x']  # returns Integer
+
+    To indicate a recursive structure, e.g. where a match should just
+    return the current type, use the symbol :attr:`.Indexable.SAME_TYPE`::
+
+        type = JSON(index_map={
+                JSON.ANY_KEY: JSON.SAME_TYPE
+            }
+        )
+
+    The above structure would cause all index expressions to keep returning
+    the same JSON structure repeatedly.   Types like the built in PostgreSQL
+    :class:`.postgresql.JSON` type use the above map by default.
+
+    For an array-like structure, the index_map would typically refer
+    to the number of dimensions known to be in the array.  E.g. if we
+    have a two-dimensional array of integers, the index_map assuming an
+    :class:`.Indexable` subclass ARRAY would be::
+
+        type = ARRAY(index_map={
+                ARRAY.ANY_KEY: {
+                    ARRAY.ANY_KEY: Integer
+                }
+            }
+        )
+
+    Above, a single-index operation returns another array::
+
+        column('x', type)[5]  # returns ARRAY of Integer
+
+    whereas calling upon two dimensions returns Integer::
+
+        column('x', type)[5][7]  # returns Integer
+
+    The above ARRAY structure is used by the built-in PostgreSQL
+    :class:`.postgresql.ARRAY` type, when configured with a specific
+    number of dimensions.
+
+    .. versionadded:: 1.1.0
+
+
     """
 
     zero_indexes = False
     """if True, Python zero-based indexes should be interpreted as one-based
     on the SQL expression side."""
 
+    ANY_KEY = util.symbol('ANY_KEY')
+    """Symbol indicating a match for any index or key in an index_map structure.
+
+    .. seealso::
+
+        :class:`.Indexable`
+
+    """
+
+    SLICE_TYPE = util.symbol('SLICE_TYPE')
+    """Symbol indicating a match for an index sent as a slice in
+    an index_map structure.
+
+    .. seealso::
+
+        :class:`.Indexable`
+
+    """
+
+    SAME_TYPE = util.symbol('SAME_TYPE')
+    """Symbol indicating that this same type should be returned upon
+    a match in an index_map structure.
+
+    .. seealso::
+
+        :class:`.Indexable`
+
+    """
+
+    index_map = util.immutabledict({
+        ANY_KEY: SAME_TYPE
+    })
+
     class Comparator(TypeEngine.Comparator):
 
-        def _index_map_type(self, right_comparator):
-            raise NotImplementedError()
+        def _type_for_index(self, index):
+            """Given a getitem index, look in the index_map to see if
+            a known type is set up for this value.
+
+            May return a new type which would contain a fragment of
+            the index map at that point.
+
+            """
+            map_ = self.type.index_map
+            if isinstance(index, slice):
+                index = Indexable.SLICE_TYPE
+
+            if index in map_:
+                new_type = map_[index]
+            else:
+                try:
+                    new_type = map_[Indexable.ANY_KEY]
+                except KeyError:
+                    raise exc.InvalidRequestError(
+                        "Key not handled by type declaration: %s" % index)
+            if isinstance(new_type, collections.Mapping):
+                return self.type.adapt(self.type.__class__, index_map=new_type)
+            elif new_type is Indexable.SAME_TYPE:
+                return self.type
+            else:
+                return new_type
+
+        def _setup_getitem(self, index):
+            return operators.getitem, index, self._type_for_index(index)
+
+        def __getitem__(self, index):
+            operator, adjusted_right_expr, result_type = \
+                self._setup_getitem(index)
+            return self.operate(
+                operator,
+                adjusted_right_expr,
+                result_type=result_type
+            )
 
         def _adapt_expression(self, op, other_comparator):
             if op is operators.getitem:
-                return op, self._index_map_type(other_comparator)
+                return op, self._type_for_index(other_comparator)
             else:
-                return super(Indexable.Comparator)._adapt_expression(
+                return super(Indexable.Comparator, self)._adapt_expression(
                     op, other_comparator)
 
     comparator_factory = Comparator
diff --git a/lib/sqlalchemy/util/langhelpers.py b/lib/sqlalchemy/util/langhelpers.py
index dd2589243..743afccfd 100644
--- a/lib/sqlalchemy/util/langhelpers.py
+++ b/lib/sqlalchemy/util/langhelpers.py
@@ -1019,7 +1019,9 @@ def constructor_copy(obj, cls, *args, **kw):
     """
 
     names = get_cls_kwargs(cls)
-    kw.update((k, obj.__dict__[k]) for k in names if k in obj.__dict__)
+    kw.update(
+        (k, obj.__dict__[k]) for k in names.difference(kw)
+        if k in obj.__dict__)
     return cls(*args, **kw)