Add sql.Operation tokenlist

3 files changed, 20 insertions, 15 deletions

diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
index 91bb3d9..3bdc303 100644
--- a/sqlparse/engine/grouping.py
+++ b/sqlparse/engine/grouping.py
@@ -85,7 +85,8 @@ def group_assignment(tlist):
 
 
 def group_comparison(tlist):
-    I_COMPERABLE = (sql.Parenthesis, sql.Function, sql.Identifier)
+    I_COMPERABLE = (sql.Parenthesis, sql.Function, sql.Identifier,
+                    sql.Operation)
     T_COMPERABLE = T_NUMERICAL + T_STRING + T_NAME
 
     func = lambda tk: imt(tk, t=T_COMPERABLE, i=I_COMPERABLE) or (
@@ -134,9 +135,9 @@ def group_arrays(tlist):
 @recurse(sql.Identifier)
 def group_operator(tlist):
     I_CYCLE = (sql.SquareBrackets, sql.Parenthesis, sql.Function,
-               sql.Identifier,)  # sql.Operation)
+               sql.Identifier, sql.Operation)
     # wilcards wouldn't have operations next to them
-    T_CYCLE = T_NUMERICAL + T_STRING + T_NAME  # + T.Wildcard
+    T_CYCLE = T_NUMERICAL + T_STRING + T_NAME
     func = lambda tk: imt(tk, i=I_CYCLE, t=T_CYCLE)
 
     token = tlist.token_next_by(t=(T.Operator, T.Wildcard))
@@ -146,8 +147,7 @@ def group_operator(tlist):
         if func(left) and func(right):
             token.ttype = T.Operator
             tokens = tlist.tokens_between(left, right)
-            # token = tlist.group_tokens(sql.Operation, tokens)
-            token = tlist.group_tokens(sql.Identifier, tokens)
+            token = tlist.group_tokens(sql.Operation, tokens)
 
         token = tlist.token_next_by(t=(T.Operator, T.Wildcard), idx=token)
 
@@ -155,7 +155,7 @@ def group_operator(tlist):
 @recurse(sql.IdentifierList)
 def group_identifier_list(tlist):
     I_IDENT_LIST = (sql.Function, sql.Case, sql.Identifier, sql.Comparison,
-                    sql.IdentifierList)  # sql.Operation
+                    sql.IdentifierList, sql.Operation)
     T_IDENT_LIST = (T_NUMERICAL + T_STRING + T_NAME +
                     (T.Keyword, T.Comment, T.Wildcard))
 
@@ -212,7 +212,7 @@ def group_where(tlist):
 @recurse()
 def group_aliased(tlist):
     I_ALIAS = (sql.Parenthesis, sql.Function, sql.Case, sql.Identifier,
-               )  # sql.Operation)
+               sql.Operation)
 
     token = tlist.token_next_by(i=I_ALIAS, t=T.Number)
     while token:
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
index 43a89e7..e24c10f 100644
--- a/sqlparse/sql.py
+++ b/sqlparse/sql.py
@@ -588,3 +588,7 @@ class Begin(TokenList):
     """A BEGIN/END block."""
     M_OPEN = T.Keyword, 'BEGIN'
     M_CLOSE = T.Keyword, 'END'
+
+
+class Operation(TokenList):
+    """Grouping of operations"""
diff --git a/tests/test_grouping.py b/tests/test_grouping.py
index 7ea1c75..147162f 100644
--- a/tests/test_grouping.py
+++ b/tests/test_grouping.py
@@ -106,15 +106,16 @@ class TestGrouping(TestCaseBase):
         self.assert_(isinstance(p.tokens[0].tokens[0], sql.Function))
         p = sqlparse.parse('foo()||col2 bar')[0]
         self.assert_(isinstance(p.tokens[0], sql.Identifier))
-        self.assert_(isinstance(p.tokens[0].tokens[0], sql.Function))
+        self.assert_(isinstance(p.tokens[0].tokens[0], sql.Operation))
+        self.assert_(isinstance(p.tokens[0].tokens[0].tokens[0], sql.Function))
 
     def test_identifier_extended(self):  # issue 15
         p = sqlparse.parse('foo+100')[0]
-        self.assert_(isinstance(p.tokens[0], sql.Identifier))
+        self.assert_(isinstance(p.tokens[0], sql.Operation))
         p = sqlparse.parse('foo + 100')[0]
-        self.assert_(isinstance(p.tokens[0], sql.Identifier))
+        self.assert_(isinstance(p.tokens[0], sql.Operation))
         p = sqlparse.parse('foo*100')[0]
-        self.assert_(isinstance(p.tokens[0], sql.Identifier))
+        self.assert_(isinstance(p.tokens[0], sql.Operation))
 
     def test_identifier_list(self):
         p = sqlparse.parse('a, b, c')[0]
@@ -267,25 +268,25 @@ class TestStatement(TestCaseBase):
 def test_identifier_with_operators():  # issue 53
     p = sqlparse.parse('foo||bar')[0]
     assert len(p.tokens) == 1
-    assert isinstance(p.tokens[0], sql.Identifier)
+    assert isinstance(p.tokens[0], sql.Operation)
     # again with whitespaces
     p = sqlparse.parse('foo || bar')[0]
     assert len(p.tokens) == 1
-    assert isinstance(p.tokens[0], sql.Identifier)
+    assert isinstance(p.tokens[0], sql.Operation)
 
 
 def test_identifier_with_op_trailing_ws():
     # make sure trailing whitespace isn't grouped with identifier
     p = sqlparse.parse('foo || bar ')[0]
     assert len(p.tokens) == 2
-    assert isinstance(p.tokens[0], sql.Identifier)
+    assert isinstance(p.tokens[0], sql.Operation)
     assert p.tokens[1].ttype is T.Whitespace
 
 
 def test_identifier_with_string_literals():
     p = sqlparse.parse('foo + \'bar\'')[0]
     assert len(p.tokens) == 1
-    assert isinstance(p.tokens[0], sql.Identifier)
+    assert isinstance(p.tokens[0], sql.Operation)
 
 
 # This test seems to be wrong. It was introduced when fixing #53, but #111