Make index expressions in flatten expression consistent

The behavior is now consistent with wildcards.

jmespath/ast.py

@@ -150,7 +150,7 @@ class WildcardIndex(AST):
 
     """
     def search(self, value):
-        return _MultiMatch(value)
+        return _Projection(value)
 
     def pretty_print(self, indent=''):
         return "%sIndex(*)" % indent
@@ -166,7 +166,7 @@ class WildcardValues(AST):
     """
     def search(self, value):
         try:
-            return _MultiMatch(value.values())
+            return _Projection(value.values())
         except AttributeError:
             return None
 
@@ -185,15 +185,15 @@ def search(self, value):
                     merged_list.extend(element)
                 else:
                     merged_list.append(element)
-            return _FlattenedMultiMatch(merged_list)
+            return _Projection(merged_list)
         else:
-            return _FlattenedMultiMatch(value)
+            return _Projection(value)
 
     def pretty_print(self, indent=''):
         return "%sListElements()" % indent
 
 
-class _BaseMultiMatch(list):
+class _Projection(list):
     def __init__(self, elements):
         self.extend(elements)
 
@@ -210,6 +210,18 @@ def get(self, value):
                 results.append(result)
         return results
 
+    def get_index(self, index):
+        matches = []
+        for el in self:
+            if not isinstance(el, list):
+                continue
+            try:
+                matches.append(el[index])
+            except (IndexError, TypeError):
+                pass
+        if matches:
+            return self.__class__(matches)
+
     def multi_get(self, nodes):
         results = self.__class__([])
         for element in self:
@@ -235,23 +247,6 @@ def multi_get_list(self, nodes):
         return results
 
 
-class _MultiMatch(_BaseMultiMatch):
-    def get_index(self, index):
-        matches = []
-        for el in self:
-            try:
-                matches.append(el[index])
-            except (IndexError, TypeError):
-                pass
-        if matches:
-            return self.__class__(matches)
-
-
-class _FlattenedMultiMatch(_BaseMultiMatch):
-    def __init__(self, elements):
-        self.extend(elements)
-
-
 class ORExpression(AST):
     def __init__(self, first, remaining):
         self.first = first

tests/compliance/indices.json

@@ -213,36 +213,20 @@
     "cases": [
         {
            "expression": "foo[][0]",
-           "result": ["one", "two"]
+           "result": ["one", "three", "five", "seven", "nine", "ten"]
         },
         {
            "expression": "foo[][1]",
-           "result": ["three", "four"]
+           "result": ["two", "four", "six", "eight"]
         },
         {
            "expression": "foo[][0][0]",
-           "result": "one"
+           "result": null
         },
-         {
-            "expression": "foo[][1][0]",
-            "result": "three"
-         },
-         {
-            "expression": "foo[][0][1]",
-            "result": "two"
-         },
-         {
-            "expression": "foo[][1][1]",
-            "result": "four"
-         },
          {
             "expression": "foo[][2][2]",
             "result": null
          },
-         {
-            "expression": "foo[][0][0][0]",
-            "result": null
-         },
          {
             "expression": "foo[][0][0][100]",
             "result": null

tests/test_parser.py

@@ -240,7 +240,8 @@ def setUp(self):
     def test_merge_with_indices(self):
         parsed = self.parser.parse('foo[][0]')
         match = parsed.search(self.data)
-        self.assertEqual(match, ["one", "two"])
+        self.assertEqual(match, ["one", "three", "five", "seven",
+                                 "nine", "ten"])
 
     def test_trailing_merged_operator(self):
         parsed = self.parser.parse('foo[]')