diff options
| -rw-r--r-- | sqlparse/engine/grouping.py | 4 | ||||
| -rw-r--r-- | sqlparse/lexer.py | 1 | ||||
| -rw-r--r-- | sqlparse/sql.py | 7 | ||||
| -rw-r--r-- | sqlparse/tokens.py | 1 | ||||
| -rw-r--r-- | tests/test_parse.py | 49 |
5 files changed, 61 insertions, 1 deletions
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 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'] + + |