diff options
| -rw-r--r-- | python/util.py | 23 |
1 files changed, 23 insertions, 0 deletions
diff --git a/python/util.py b/python/util.py index 5cde278..9339a8e 100644 --- a/python/util.py +++ b/python/util.py @@ -28,6 +28,13 @@ def get_tokens(code): except TokenError: return [] +def all_tokens(code): + try: + stream = io.BytesIO(code.encode('utf-8')) + return [t for t in tokenize(stream.readline)] + except TokenError: + return [] + # Check if tokens contain a sequence of tokens (given as a list of strings). def has_token_sequence(tokens, sequence): for i in range(len(tokens)-len(sequence)+1): @@ -85,6 +92,22 @@ def get_exception_desc(exc): return [{'id':'error', 'args': {'message': exc}}] return None + +def has_comprehension(code): + """ Searches code for comprehensions and generators. """ + for l in code.split('\n'): + if 'for' in l and ':' not in l: + 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): + return True + return False + if __name__ == '__main__': print(has_token_sequence(get_tokens('x + y >= 0'), ['>=', '0'])) print(has_token_sequence(get_tokens('x + y > 0'), ['>=', '0'])) |