summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--python/util.py23
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