Group square-brackets into identifiers

darikg · darikg · commit b61fe36f718c · 2015-03-04T10:41:34.000-05:00
Indentifier.get_array_indices() looks for square brackets, and yields lists of bracket grouped tokens as array indices
diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
@@ -159,16 +159,17 @@ def _consume_cycle(tl, i):
             lambda y: (y.match(T.Punctuation, '.')
                        or y.ttype in (T.Operator,
                                       T.Wildcard,
-                                      T.ArrayIndex,
-                                      T.Name)),
+                                      T.Name)
+                       or isinstance(y, sql.SquareBrackets)),
             lambda y: (y.ttype in (T.String.Symbol,
                                    T.Name,
                                    T.Wildcard,
-                                   T.ArrayIndex,
                                    T.Literal.String.Single,
                                    T.Literal.Number.Integer,
                                    T.Literal.Number.Float)
-                       or isinstance(y, (sql.Parenthesis, sql.Function)))))
+                       or isinstance(y, (sql.Parenthesis,
+                                         sql.SquareBrackets,
+                                         sql.Function)))))
         for t in tl.tokens[i:]:
             # Don't take whitespaces into account.
             if t.ttype is T.Whitespace:
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
@@ -511,11 +511,12 @@ def get_ordering(self):
         return ordering.value.upper()
 
     def get_array_indices(self):
-        """Returns an iterator of index expressions as strings"""
+        """Returns an iterator of index token lists"""
 
-        # Use [1:-1] index to discard the square brackets
-        return (tok.value[1:-1] for tok in self.tokens
-                if tok.ttype in T.ArrayIndex)
+        for tok in self.tokens:
+            if isinstance(tok, SquareBrackets):
+                # Use [1:-1] index to discard the square brackets
+                yield tok.tokens[1:-1]
 
 
 class IdentifierList(TokenList):
