From 7a8ad0abe660bc05ef3b1e48015758322a0d971c Mon Sep 17 00:00:00 2001
From: Dhruv Manilawala <dhruvmanila@gmail.com>
Date: Tue, 14 May 2024 08:44:52 +0530
Subject: [PATCH] Use `TokenKind` in `doc_lines_from_tokens`

---
 crates/ruff_linter/src/doc_lines.rs | 19 +++++++++----------
 crates/ruff_linter/src/linter.rs    |  2 +-
 2 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/crates/ruff_linter/src/doc_lines.rs b/crates/ruff_linter/src/doc_lines.rs
index 7b69fb866920d1..eebe21cec3ca54 100644
--- a/crates/ruff_linter/src/doc_lines.rs
+++ b/crates/ruff_linter/src/doc_lines.rs
@@ -4,27 +4,26 @@
 use std::iter::FusedIterator;
 
 use ruff_python_ast::{self as ast, Stmt, Suite};
-use ruff_python_parser::lexer::LexResult;
-use ruff_python_parser::Tok;
+use ruff_python_parser::{TokenKind, TokenKindIter};
 use ruff_text_size::{Ranged, TextSize};
 
 use ruff_python_ast::statement_visitor::{walk_stmt, StatementVisitor};
 use ruff_source_file::{Locator, UniversalNewlineIterator};
 
 /// Extract doc lines (standalone comments) from a token sequence.
-pub(crate) fn doc_lines_from_tokens(lxr: &[LexResult]) -> DocLines {
-    DocLines::new(lxr)
+pub(crate) fn doc_lines_from_tokens(tokens: TokenKindIter) -> DocLines {
+    DocLines::new(tokens)
 }
 
 pub(crate) struct DocLines<'a> {
-    inner: std::iter::Flatten<core::slice::Iter<'a, LexResult>>,
+    inner: TokenKindIter<'a>,
     prev: TextSize,
 }
 
 impl<'a> DocLines<'a> {
-    fn new(lxr: &'a [LexResult]) -> Self {
+    fn new(tokens: TokenKindIter<'a>) -> Self {
         Self {
-            inner: lxr.iter().flatten(),
+            inner: tokens,
             prev: TextSize::default(),
         }
     }
@@ -39,15 +38,15 @@ impl Iterator for DocLines<'_> {
             let (tok, range) = self.inner.next()?;
 
             match tok {
-                Tok::Comment(..) => {
+                TokenKind::Comment => {
                     if at_start_of_line {
                         break Some(range.start());
                     }
                 }
-                Tok::Newline | Tok::NonLogicalNewline => {
+                TokenKind::Newline | TokenKind::NonLogicalNewline => {
                     at_start_of_line = true;
                 }
-                Tok::Indent | Tok::Dedent => {
+                TokenKind::Indent | TokenKind::Dedent => {
                     // ignore
                 }
                 _ => {
diff --git a/crates/ruff_linter/src/linter.rs b/crates/ruff_linter/src/linter.rs
index 7033dbab97567f..c8bf5656f5f5fe 100644
--- a/crates/ruff_linter/src/linter.rs
+++ b/crates/ruff_linter/src/linter.rs
@@ -93,7 +93,7 @@ pub fn check_path(
     let use_doc_lines = settings.rules.enabled(Rule::DocLineTooLong);
     let mut doc_lines = vec![];
     if use_doc_lines {
-        doc_lines.extend(doc_lines_from_tokens(&tokens));
+        doc_lines.extend(doc_lines_from_tokens(tokens.kinds()));
     }
 
     // Run the token-based rules.