Added python util functions.

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']))