Fix token-parent behavior

Closes issue #226

2 files changed, 11 insertions, 0 deletions

diff --git a/sqlparse/sql.py b/sqlparse/sql.py
index cee6af5..52b3bf1 100644
--- a/sqlparse/sql.py
+++ b/sqlparse/sql.py
@@ -140,6 +140,7 @@ class TokenList(Token):
 
     def __init__(self, tokens=None):
         self.tokens = tokens or []
+        [setattr(token, 'parent', self) for token in tokens]
         super(TokenList, self).__init__(None, text_type(self))
 
     def __str__(self):
@@ -300,11 +301,13 @@ class TokenList(Token):
 
     def insert_before(self, where, token):
         """Inserts *token* before *where*."""
+        token.parent = self
         self.tokens.insert(self.token_index(where), token)
 
     def insert_after(self, where, token, skip_ws=True):
         """Inserts *token* after *where*."""
         next_token = self.token_next(where, skip_ws=skip_ws)
+        token.parent = self
         if next_token is None:
             self.tokens.append(token)
         else:
diff --git a/tests/test_parse.py b/tests/test_parse.py
index 0f5af04..75a7ab5 100644
--- a/tests/test_parse.py
+++ b/tests/test_parse.py
@@ -390,3 +390,11 @@ def test_wildcard_multiplication():
 
     p = sqlparse.parse('select 1 * 2 from dual')[0]
     assert p.tokens[2][2].ttype == T.Operator
+
+
+def test_stmt_tokens_parents():
+    # see issue 226
+    sql = "CREATE TABLE test();"
+    stmt = sqlparse.parse(sql)[0]
+    for token in stmt.tokens:
+        assert token.has_ancestor(stmt)