From e1b5482b1413591260f41ceb50e97091928f195a Mon Sep 17 00:00:00 2001
From: Andi Albrecht <albrecht.andi@gmail.com>
Date: Sun, 8 Feb 2015 16:52:15 +0100
Subject: (postgresql) Add support for square bracket array indexing, by
 darikg.

Pull request #170 with trivial conflicts resolved.
---
 sqlparse/engine/grouping.py | 4 +++-
 sqlparse/lexer.py           | 1 +
 sqlparse/sql.py             | 7 +++++++
 sqlparse/tokens.py          | 1 +
 4 files changed, 12 insertions(+), 1 deletion(-)

(limited to 'sqlparse')

diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
index d6f1360..47d3c2e 100644
--- a/sqlparse/engine/grouping.py
+++ b/sqlparse/engine/grouping.py
@@ -156,10 +156,12 @@ def group_identifier(tlist):
         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
index 611835f..a60c789 100644
--- 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
index 6174db0..c974b35 100644
--- a/sqlparse/sql.py
+++ b/sqlparse/sql.py
@@ -502,6 +502,13 @@ class Identifier(TokenList):
             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
index 01a9b89..014984b 100644
--- a/sqlparse/tokens.py
+++ b/sqlparse/tokens.py
@@ -57,6 +57,7 @@ Literal = Token.Literal
 String = Literal.String
 Number = Literal.Number
 Punctuation = Token.Punctuation
+ArrayIndex = Punctuation.ArrayIndex
 Operator = Token.Operator
 Comparison = Operator.Comparison
 Wildcard = Token.Wildcard
-- 
cgit v1.2.1