Modify current count routine to only work for Feature.token

Signed-off-by: Matthew Peveler <matt.peveler@gmail.com>

Author: MasterOdin

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) {

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