1 files changed, 15 insertions, 8 deletions
diff --git a/python/util.py b/python/util.py
index 9339a8e..f04fcc2 100644
--- a/python/util.py
+++ b/python/util.py
@@ -19,6 +19,7 @@
 import io
 import re
 from tokenize import tokenize, TokenError
+import ast
 
 def get_tokens(code):
     """ Gets a list of tokens. """
@@ -92,19 +93,25 @@ def get_exception_desc(exc):
             return [{'id':'error', 'args': {'message': exc}}]
     return None
 
+def get_ast(code):
+    """
+    Turn code into ast; use it when regular expressions on strings
+    are not enjoyable enough.
+
+    """
+    return ast.parse(code)
 
-def has_comprehension(code):
+def has_comprehension(tree):
     """  Searches code for comprehensions and generators. """
-    for l in code.split('\n'):
-        if 'for' in l and ':' not in l:
+    for n in ast.walk(tree):
+        if isinstance(n, ast.comprehension):
             return True
     return False
 
-def has_loop(code):
-    """  Searches code for loops (for and while), but skips comprehensions and
-     generators. """
-    for l in code.split('\n'):
-        if ('for' in l and ':' in l) or ('while' in l and ':' in l):
+def has_loop(tree):
+    """  Searches abstract syntax tree for loops (for and while). """
+    for n in ast.walk(tree):
+        if isinstance(n, ast.For) or isinstance(n, ast.While):
             return True
     return False