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

Pull request #170 with trivial conflicts resolved.

1 files changed, 49 insertions, 0 deletions

diff --git a/tests/test_parse.py b/tests/test_parse.py
index 2650d7f..24eea2d 100644
--- 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']
+
+