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