token_next shouldn't ignore skip_cm

darikg · darikg · commit 66ae504c4bf8 · 2016-06-20T09:56:42.000-04:00
diff --git a/sqlparse/sql.py b/sqlparse/sql.py
@@ -257,37 +257,25 @@ def token_prev(self, idx, skip_ws=True, skip_cm=False):
         """Returns the previous token relative to *idx*.
 
         If *skip_ws* is ``True`` (the default) whitespace tokens are ignored.
+        If *skip_cm* is ``True`` comments are ignored.
         ``None`` is returned if there's no previous token.
         """
-        if idx is None:
-            return None, None
-        idx += 1  # alot of code usage current pre-compensates for this
-        funcs = lambda tk: not ((skip_ws and tk.is_whitespace()) or
-                                (skip_cm and imt(tk, t=T.Comment, i=Comment)))
-        return self._token_matching(funcs, idx, reverse=True)
+        return self.token_next(idx, skip_ws, skip_cm, _reverse=True)
 
-    # TODO: May need to implement skip_cm for upstream changes.
     # TODO: May need to re-add default value to idx
-    def token_next(self, idx, skip_ws=True, skip_cm=False):
+    def token_next(self, idx, skip_ws=True, skip_cm=False, _reverse=False):
         """Returns the next token relative to *idx*.
 
         If *skip_ws* is ``True`` (the default) whitespace tokens are ignored.
+        If *skip_cm* is ``True`` comments are ignored.
         ``None`` is returned if there's no next token.
         """
         if idx is None:
             return None, None
         idx += 1  # alot of code usage current pre-compensates for this
-        try:
-            if not skip_ws:
-                return idx, self.tokens[idx]
-            else:
-                while True:
-                    token = self.tokens[idx]
-                    if not token.is_whitespace():
-                        return idx, token
-                    idx += 1
-        except IndexError:
-            return None, None
+        funcs = lambda tk: not ((skip_ws and tk.is_whitespace()) or
+                                (skip_cm and imt(tk, t=T.Comment, i=Comment)))
+        return self._token_matching(funcs, idx, reverse=_reverse)
 
     def token_index(self, token, start=0):
         """Return list index of token."""
diff --git a/tests/test_regressions.py b/tests/test_regressions.py
@@ -312,3 +312,9 @@ def test_issue207_runaway_format():
                            "          2 as two,",
                            "          3",
                            "   from dual) t0"])
+
+
+def token_next_doesnt_ignore_skip_cm():
+    sql = '--comment\nselect 1'
+    tok = sqlparse.parse(sql)[0].token_next(-1, skip_cm=True)[1]
+    assert tok.value == 'select'
