Modify current count routine to only work for Feature.token

Signed-off-by: Matthew Peveler <[email protected]>
Submitty · Feb 22, 2022 · 1d96305 · 1d96305
1 parent 2cd0f0c
commit 1d96305
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/src/count.ts b/src/count.ts
@@ -28,7 +28,10 @@ export function countToken(tree: Tree, feature: Feature, token: string): number
   const cursor = tree.rootNode.walk();
   while (true) {
     if (cursor.gotoFirstChild() || cursor.gotoNextSibling()) {
-      if (convertTokenName(cursor.nodeType).toLowerCase() === tokenLower) {
+      if (
+        feature === Feature.token &&
+        convertTokenName(cursor.nodeType).toLowerCase() === tokenLower
+      ) {
         count++;
       }
       continue;
@@ -41,7 +44,11 @@ export function countToken(tree: Tree, feature: Feature, token: string): number
         break;
       }
     }
-    if (hadSibling && convertTokenName(cursor.nodeType).toLowerCase() === tokenLower) {
+    if (
+      feature === Feature.token &&
+      hadSibling &&
+      convertTokenName(cursor.nodeType).toLowerCase() === tokenLower
+    ) {
       count++;
     }
     if (cursor.currentNode === tree.rootNode) {

diff --git a/test/diagnostics.spec.ts b/test/diagnostics.spec.ts
@@ -63,6 +63,8 @@ describe('diagnostics', () => {
       });
     });
 
+    // Testing that we find tokens that are at end of lines, after tokens with
+    // children
     test('find all colon tokens', () => {
       const { tokens } = diagnostics(
         parseFile(Language.python, path.join(TEST_DATA, 'python', 'operators.py')),