diff options
author | Martin Možina <martin.mozina@fri.uni-lj.si> | 2016-09-27 16:35:52 +0200 |
---|---|---|
committer | Martin Možina <martin.mozina@fri.uni-lj.si> | 2016-09-27 16:35:52 +0200 |
commit | 2013b5b59f2400ac3d4bf218a7ec5b1174f812d7 (patch) | |
tree | e1db4868c70ad92bbc5170ecc0f588d3dc06cf23 /python | |
parent | 4e2ea9ab2da792bae138916f45e4cf33ffca6ff0 (diff) |
Added ast related functions.
Diffstat (limited to 'python')
-rw-r--r-- | python/util.py | 23 |
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 |