(postgresql) Add support for square bracket array indexing, by darikg.

andialbrecht · andialbrecht · commit e1b5482b1413 · 2015-02-08T16:52:15.000+01:00
Pull request andialbrecht#170 with trivial conflicts resolved.
diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
@@ -156,10 +156,12 @@ def _consume_cycle(tl, i):
         x = itertools.cycle((
             lambda y: (y.match(T.Punctuation, '.')
                        or y.ttype is T.Operator
-                       or y.ttype is T.Wildcard),
+                       or y.ttype is T.Wildcard
+                       or y.ttype is T.ArrayIndex),
             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)
diff --git a/sqlparse/lexer.py b/sqlparse/lexer.py
@@ -195,6 +195,7 @@ class Lexer(object):
             (r"'(''|\\'|[^'])*'", tokens.String.Single),
             # not a real string literal in ANSI SQL:
             (r'(""|".*?[^\\]")', tokens.String.Symbol),
+            (r'(?<=[\w\]])(\[[^\]]*?\])', tokens.Punctuation.ArrayIndex),
             (r'(\[[^\]]+\])', tokens.Name),
             (r'((LEFT\s+|RIGHT\s+|FULL\s+)?(INNER\s+|OUTER\s+|STRAIGHT\s+)?|(CROSS\s+|NATURAL\s+)?)?JOIN\b', tokens.Keyword),
             (r'END(\s+IF|\s+LOOP)?\b', tokens.Keyword),
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
@@ -502,6 +502,13 @@ def get_ordering(self):
             return None
         return ordering.value.upper()
 
+    def get_array_indices(self):
+        """Returns an iterator of index expressions as strings"""
+
+        # 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)
+
 
 class IdentifierList(TokenList):
     """A list of :class:`~sqlparse.sql.Identifier`\'s."""
diff --git a/sqlparse/tokens.py b/sqlparse/tokens.py
@@ -57,6 +57,7 @@ def __repr__(self):
 String = Literal.String
 Number = Literal.Number
 Punctuation = Token.Punctuation
+ArrayIndex = Punctuation.ArrayIndex
 Operator = Token.Operator
 Comparison = Operator.Comparison
 Wildcard = Token.Wildcard
diff --git a/tests/test_parse.py b/tests/test_parse.py
@@ -205,3 +205,52 @@ def test_single_quotes_with_linebreaks():  # issue118
     assert len(p) == 1
     assert p[0].ttype is T.String.Single
 
+
+def test_array_indexed_column():
+    # Make sure we still parse sqlite style escapes
+    p = sqlparse.parse('[col1],[col2]')[0].tokens
+    assert (len(p) == 1
+            and isinstance(p[0], sqlparse.sql.IdentifierList)
+            and [id.get_name() for id in p[0].get_identifiers()]
+                    == ['[col1]', '[col2]'])
+
+    p = sqlparse.parse('[col1]+[col2]')[0]
+    types = [tok.ttype for tok in p.flatten()]
+    assert types == [T.Name, T.Operator, T.Name]
+
+    p = sqlparse.parse('col[1]')[0].tokens
+    assert (len(p) == 1
+        and tuple(p[0].get_array_indices()) == ('1',)
+        and p[0].get_name() == 'col')
+
+    p = sqlparse.parse('col[1][1:5] as mycol')[0].tokens
+    assert (len(p) == 1
+        and tuple(p[0].get_array_indices()) == ('1', '1:5')
+        and p[0].get_name() == 'mycol'
+        and p[0].get_real_name() == 'col')
+
+    p = sqlparse.parse('col[1][other_col]')[0].tokens
+    assert len(p) == 1 and tuple(p[0].get_array_indices()) == ('1', 'other_col')
+
+    sql = 'SELECT col1, my_1d_array[2] as alias1, my_2d_array[2][5] as alias2'
+    p = sqlparse.parse(sql)[0].tokens
+    assert len(p) == 3 and isinstance(p[2], sqlparse.sql.IdentifierList)
+    ids = list(p[2].get_identifiers())
+    assert (ids[0].get_name() == 'col1'
+            and tuple(ids[0].get_array_indices()) == ()
+            and ids[1].get_name() == 'alias1'
+            and ids[1].get_real_name() == 'my_1d_array'
+            and tuple(ids[1].get_array_indices()) == ('2',)
+            and ids[2].get_name() == 'alias2'
+            and ids[2].get_real_name() == 'my_2d_array'
+            and tuple(ids[2].get_array_indices()) == ('2', '5'))
+
+
+def test_typed_array_definition():
+    # array indices aren't grouped with builtins, but make sure we can extract
+    # indentifer names
+    p = sqlparse.parse('x int, y int[], z int')[0]
+    names = [x.get_name() for x in p.get_sublists()]
+    assert names == ['x', 'y', 'z']
+
+
