Speed up has_print_function.

baron/future.py

@@ -2,38 +2,35 @@
 
 
 def has_print_function(tokens):
-    for pos in range(len(tokens)):
-        if tokens_define_print_function(tokens[pos:]):
-            return True
-    return False
-
-
-def tokens_define_print_function(tokens):
-    token = iter(tokens)
-    try:
-        if next(token)[0] != 'FROM':
-            return False
-        if next(token)[0:2] != ('NAME', '__future__'):
-            return False
-        if next(token)[0] != 'IMPORT':
-            return False
-
-        current_token = next(token)
+    p = 0
+    while p < len(tokens):
+        if tokens[p][0] != 'FROM':
+            p += 1
+            continue
+        if tokens[p + 1][0:2] != ('NAME', '__future__'):
+            p += 1
+            continue
+        if tokens[p + 2][0] != 'IMPORT':
+            p += 1
+            continue
+
+        current = p + 3
         # ignore LEFT_PARENTHESIS token
-        if current_token[0] == 'LEFT_PARENTHESIS':
-            current_token = next(token)
+        if tokens[current][0] == 'LEFT_PARENTHESIS':
+            current += 1
 
-        while (current_token[0] == 'NAME'):
-            if current_token[1] == 'print_function':
+        while (current < len(tokens) and tokens[current][0] == 'NAME'):
+            if tokens[current][1] == 'print_function':
                 return True
+
             # ignore AS and NAME tokens if present
             # anyway, ignore COMMA token
-            if next(token)[0] == 'AS':
-                next(token)
-                next(token)
-            current_token = next(token)
-    except StopIteration:
-        pass
+            if current + 1 < len(tokens) and tokens[current + 1][0] == 'AS':
+                current += 4
+            else:
+                current += 2
+        p += 1
+
     return False
 
 
@@ -42,4 +39,3 @@ def is_print(token):
         return token[0] == 'PRINT'
 
     return [('NAME', 'print') if is_print(x) else x for x in tokens]
-