Fix column_defs example

1 files changed, 18 insertions, 22 deletions

diff --git a/examples/column_defs_lowlevel.py b/examples/column_defs_lowlevel.py
index 5acbdec..5e98be3 100644
--- a/examples/column_defs_lowlevel.py
+++ b/examples/column_defs_lowlevel.py
@@ -11,43 +11,39 @@
 
 import sqlparse
 
-SQL = """CREATE TABLE foo (
-  id integer primary key,
-  title varchar(200) not null,
-  description text
-);"""
-
-
-parsed = sqlparse.parse(SQL)[0]
-
-# extract the parenthesis which holds column definitions
-par = parsed.token_next_by(i=sqlparse.sql.Parenthesis)
-
 
 def extract_definitions(token_list):
     # assumes that token_list is a parenthesis
     definitions = []
     tmp = []
-    # grab the first token, ignoring whitespace
-    token = token_list.token_next(0)
+    # grab the first token, ignoring whitespace. idx=1 to skip open (
+    token = token_list.token_next(1)
     while token and not token.match(sqlparse.tokens.Punctuation, ')'):
         tmp.append(token)
-        idx = token_list.token_index(token)
         # grab the next token, this times including whitespace
-        token = token_list.token_next(idx, skip_ws=False)
+        token = token_list.token_next(token, skip_ws=False)
         # split on ",", except when on end of statement
         if token and token.match(sqlparse.tokens.Punctuation, ','):
             definitions.append(tmp)
             tmp = []
-            idx = token_list.token_index(token)
-            token = token_list.token_next(idx)
+            token = token_list.token_next(token)
     if tmp and isinstance(tmp[0], sqlparse.sql.Identifier):
         definitions.append(tmp)
     return definitions
 
 
-columns = extract_definitions(par)
+if __name__ == '__main__':
+    SQL = """CREATE TABLE foo (
+             id integer primary key,
+             title varchar(200) not null,
+             description text);"""
+
+    parsed = sqlparse.parse(SQL)[0]
+
+    # extract the parenthesis which holds column definitions
+    par = parsed.token_next_by(i=sqlparse.sql.Parenthesis)
+    columns = extract_definitions(par)
 
-for column in columns:
-    print('NAME: {name:10} DEFINITION: {definition}'.format(
-        name=column[0], definition=''.join(str(t) for t in column[1:])))
+    for column in columns:
+        print('NAME: {name:10} DEFINITION: {definition}'.format(
+            name=column[0], definition=''.join(str(t) for t in column[1:])))