Support for BEGIN/END blocks, add FOREACH keyword.

5 files changed, 62 insertions, 2 deletions

diff --git a/CHANGES b/CHANGES
index d29473a..1291ae5 100644
--- a/CHANGES
+++ b/CHANGES
@@ -6,6 +6,8 @@ Bug Fixes
 
 Enhancements
 * Classify DML keywords (issue116, by Victor Hahn).
+* Add missing FOREACH keyword.
+* Grouping of BEGIN/END blocks.
 
 
 Release 0.1.10 (Nov 02, 2013)
diff --git a/sqlparse/engine/grouping.py b/sqlparse/engine/grouping.py
index 47e77ac..29f7a37 100644
--- a/sqlparse/engine/grouping.py
+++ b/sqlparse/engine/grouping.py
@@ -99,6 +99,16 @@ def group_for(tlist):
                     sql.For, True)
 
 
+def group_foreach(tlist):
+    _group_matching(tlist, T.Keyword, 'FOREACH', T.Keyword, 'END LOOP',
+                    sql.For, True)
+
+
+def group_begin(tlist):
+    _group_matching(tlist, T.Keyword, 'BEGIN', T.Keyword, 'END',
+                    sql.Begin, True)
+
+
 def group_as(tlist):
 
     def _right_valid(token):
@@ -369,5 +379,7 @@ def group(tlist):
             group_comparison,
             group_identifier_list,
             group_if,
-            group_for]:
+            group_for,
+            group_foreach,
+            group_begin]:
         func(tlist)
diff --git a/sqlparse/keywords.py b/sqlparse/keywords.py
index f1d76d0..1675d81 100644
--- a/sqlparse/keywords.py
+++ b/sqlparse/keywords.py
@@ -29,7 +29,7 @@ KEYWORDS = {
 
     'BACKWARD': tokens.Keyword,
     'BEFORE': tokens.Keyword,
-    'BEGIN': tokens.Keyword.DML,
+    'BEGIN': tokens.Keyword,
     'BETWEEN': tokens.Keyword,
     'BITVAR': tokens.Keyword,
     'BIT_LENGTH': tokens.Keyword,
@@ -163,6 +163,7 @@ KEYWORDS = {
     'FINAL': tokens.Keyword,
     'FIRST': tokens.Keyword,
     'FORCE': tokens.Keyword,
+    'FOREACH': tokens.Keyword,
     'FOREIGN': tokens.Keyword,
     'FORTRAN': tokens.Keyword,
     'FORWARD': tokens.Keyword,
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
index 9567aa5..cfd622b 100644
--- a/sqlparse/sql.py
+++ b/sqlparse/sql.py
@@ -629,3 +629,9 @@ class Function(TokenList):
             elif isinstance(t, Identifier):
                 return [t,]
         return []
+
+
+class Begin(TokenList):
+    """A BEGIN/END block."""
+
+    __slots__ = ('value', 'ttype', 'tokens')
diff --git a/tests/test_grouping.py b/tests/test_grouping.py
index f4bfb1a..24b50dc 100644
--- a/tests/test_grouping.py
+++ b/tests/test_grouping.py
@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 
+import pytest
+
 import sqlparse
 from sqlparse import sql
 from sqlparse import tokens as T
@@ -277,3 +279,40 @@ def test_comparison_with_parenthesis():  # issue23
     comp = p.tokens[0]
     assert isinstance(comp.left, sql.Parenthesis)
     assert comp.right.ttype is T.Number.Integer
+
+
+@pytest.mark.parametrize('start', ['FOR', 'FOREACH'])
+def test_forloops(start):
+    p = sqlparse.parse('%s foo in bar LOOP foobar END LOOP' % start)[0]
+    assert (len(p.tokens)) == 1
+    assert isinstance(p.tokens[0], sql.For)
+
+
+def test_nested_for():
+    p = sqlparse.parse('FOR foo LOOP FOR bar LOOP END LOOP END LOOP')[0]
+    assert len(p.tokens) == 1
+    for1 = p.tokens[0]
+    assert for1.tokens[0].value == 'FOR'
+    assert for1.tokens[-1].value == 'END LOOP'
+    for2 = for1.tokens[6]
+    assert isinstance(for2, sql.For)
+    assert for2.tokens[0].value == 'FOR'
+    assert for2.tokens[-1].value == 'END LOOP'
+
+
+def test_begin():
+    p = sqlparse.parse('BEGIN foo END')[0]
+    assert len(p.tokens) == 1
+    assert isinstance(p.tokens[0], sql.Begin)
+
+
+def test_nested_begin():
+    p = sqlparse.parse('BEGIN foo BEGIN bar END END')[0]
+    assert len(p.tokens) == 1
+    outer = p.tokens[0]
+    assert outer.tokens[0].value == 'BEGIN'
+    assert outer.tokens[-1].value == 'END'
+    inner = outer.tokens[4]
+    assert inner.tokens[0].value == 'BEGIN'
+    assert inner.tokens[-1].value == 'END'
+    assert isinstance(inner, sql.Begin)