From 3c860e20748ba1e451182125ba37224473bbf9cd Mon Sep 17 00:00:00 2001
From: Omar Valdez <omarantoniovaldezf2@gmail.com>
Date: Sun, 31 Mar 2024 19:51:20 -0700
Subject: [PATCH] Fix `Unknown` and `Any` variables being coloured as Unbound

Fixes #180
---
 .../src/analyzer/semanticTokensWalker.ts      |  7 +++++-
 .../samples/semantic_highlighting/unknown.py  |  8 +++++++
 .../src/tests/semanticTokensProvider.test.ts  | 22 +++++++++++++++++++
 3 files changed, 36 insertions(+), 1 deletion(-)
 create mode 100644 packages/pyright-internal/src/tests/samples/semantic_highlighting/unknown.py

diff --git a/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts b/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts
index b6be6280d..83e8d6524 100644
--- a/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts
+++ b/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts
@@ -160,6 +160,7 @@ export class SemanticTokensWalker extends ParseTreeWalker {
                 this._addItem(node.start, node.length, SemanticTokenTypes.namespace, []);
                 return;
             // handled bellow
+            case TypeCategory.Any:
             case TypeCategory.Unknown:
             case TypeCategory.TypeVar:
                 break;
@@ -218,7 +219,11 @@ export class SemanticTokensWalker extends ParseTreeWalker {
             // with @classmethod decorator, `__new__`, `__init_subclass__`, etc.) so we need to
             // check first if it's a parameter before checking that it's a TypeVar
             this._addItem(node.start, node.length, SemanticTokenTypes.typeParameter, []);
-        } else if (type?.category === TypeCategory.Unknown) {
+            return;
+        } else if (
+            (type?.category === TypeCategory.Unknown || type?.category === TypeCategory.Any) &&
+            (declarations === undefined || declarations.length === 0)
+        ) {
             return;
         } else if (isConstantName(node.value) || (symbol && this._evaluator.isFinalVariable(symbol))) {
             this._addItem(node.start, node.length, SemanticTokenTypes.variable, [SemanticTokenModifiers.readonly]);
diff --git a/packages/pyright-internal/src/tests/samples/semantic_highlighting/unknown.py b/packages/pyright-internal/src/tests/samples/semantic_highlighting/unknown.py
new file mode 100644
index 000000000..0ad0af415
--- /dev/null
+++ b/packages/pyright-internal/src/tests/samples/semantic_highlighting/unknown.py
@@ -0,0 +1,8 @@
+from typing import Any
+
+def f(l: list) -> Any:
+    v = l[0]
+    return v
+
+g(foo)
+bar = f(...)
diff --git a/packages/pyright-internal/src/tests/semanticTokensProvider.test.ts b/packages/pyright-internal/src/tests/semanticTokensProvider.test.ts
index f0685c911..7e378929c 100644
--- a/packages/pyright-internal/src/tests/semanticTokensProvider.test.ts
+++ b/packages/pyright-internal/src/tests/semanticTokensProvider.test.ts
@@ -220,8 +220,10 @@ if (process.platform !== 'win32' || !process.env['CI']) {
             { type: 'parameter', modifiers: ['definition'], start: 140, length: 1 }, // x
             { type: 'parameter', modifiers: [], start: 159, length: 1 }, // x
             { type: 'parameter', modifiers: [], start: 163, length: 1 }, // y
+            { type: 'variable', modifiers: [], start: 169, length: 1 }, // z
             { type: 'parameter', modifiers: [], start: 177, length: 1 }, // x
             { type: 'function', modifiers: [], start: 190, length: 1 }, // g
+            { type: 'variable', modifiers: [], start: 192, length: 1 }, // z
             // lambda
             { type: 'parameter', modifiers: ['definition'], start: 203, length: 1 }, // a
             { type: 'parameter', modifiers: ['definition'], start: 206, length: 1 }, // b
@@ -229,6 +231,26 @@ if (process.platform !== 'win32' || !process.env['CI']) {
             { type: 'parameter', modifiers: [], start: 213, length: 1 }, // b
         ]);
     });
+
+    test('Unknown and Any', () => {
+        const result = semanticTokenizeSampleFile('unknown.py');
+        expect(result).toStrictEqual([
+            { type: 'namespace', modifiers: [], start: 5, length: 6 }, // typing
+            { type: 'variable', modifiers: [], start: 19, length: 3 }, // Any
+            { type: 'variable', modifiers: [], start: 19, length: 3 },
+            { type: 'function', modifiers: ['definition'], start: 28, length: 1 }, // f
+            { type: 'function', modifiers: [], start: 28, length: 1 },
+            { type: 'parameter', modifiers: ['definition'], start: 30, length: 1 }, // l
+            { type: 'class', modifiers: ['defaultLibrary', 'builtin'], start: 33, length: 4 }, // list
+            { type: 'variable', modifiers: [], start: 42, length: 3 }, // Any
+            { type: 'variable', modifiers: [], start: 51, length: 1 }, // v
+            { type: 'parameter', modifiers: [], start: 55, length: 1 }, // l
+            { type: 'variable', modifiers: [], start: 71, length: 1 }, // v
+            // `g` and `foo` should be ignored
+            { type: 'variable', modifiers: [], start: 81, length: 3 }, // bar
+            { type: 'function', modifiers: [], start: 87, length: 1 }, // f
+        ]);
+    });
 } else {
     // prevent jest from failing because no tests were found
     test('windows placeholder', () => {});