From 2013b5b59f2400ac3d4bf218a7ec5b1174f812d7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Martin=20Mo=C5=BEina?= Date: Tue, 27 Sep 2016 16:35:52 +0200 Subject: Added ast related functions. --- python/util.py | 23 +++++++++++++++-------- 1 file 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 -- cgit v1.2.1