From 9e0960bb235ae2b314119d32ebc9f636956c2073 Mon Sep 17 00:00:00 2001
From: Jurica Bradaric <jbradaric@gmail.com>
Date: Fri, 16 Jun 2023 19:43:14 +0200
Subject: [PATCH] Generate semantic tokens for modules

---
 .../src/analyzer/semanticTokensWalker.ts      | 24 ++++++++++++++++++-
 .../languageService/semanticTokensProvider.ts |  1 +
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts b/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts
index bb1de7d4d682..ec98ec6f11dd 100644
--- a/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts
+++ b/packages/pyright-internal/src/analyzer/semanticTokensWalker.ts
@@ -1,7 +1,7 @@
 import { ParseTreeWalker } from './parseTreeWalker';
 import { TypeEvaluator } from './typeEvaluatorTypes';
 import { FunctionType, OverloadedFunctionType, TypeCategory, TypeFlags } from './types';
-import { ClassNode, FunctionNode, MemberAccessNode } from '../parser/parseNodes';
+import { ClassNode, FunctionNode, ImportAsNode, ImportFromNode, ImportNode, MemberAccessNode, ModuleNameNode, ModuleNode } from '../parser/parseNodes';
 import { SemanticTokenModifiers, SemanticTokenTypes } from 'vscode-languageserver';
 
 type SemanticTokenItem = {
@@ -90,6 +90,28 @@ export class SemanticTokensWalker extends ParseTreeWalker {
                 this._addItem(node.memberName.start, node.memberName.length, SemanticTokenTypes.property, []);
                 break;
         }
+
+        const exprType = this._evaluator?.getType(node.leftExpression);
+        if (exprType?.category === TypeCategory.Module) {
+            this._addItem(node.leftExpression.start, node.leftExpression.length, SemanticTokenTypes.namespace, []);
+        }
         return super.visitMemberAccess(node);
     }
+
+    override visitImportAs(node: ImportAsNode): boolean {
+        for (const part of node.module.nameParts) {
+            this._addItem(part.start, part.length, SemanticTokenTypes.namespace, []);
+        }
+        if (node.alias) {
+            this._addItem(node.alias.start, node.alias.length, SemanticTokenTypes.namespace, []);
+        }
+        return super.visitImportAs(node);
+    }
+
+    override visitImportFrom(node: ImportFromNode): boolean {
+        for (const part of node.module.nameParts) {
+            this._addItem(part.start, part.length, SemanticTokenTypes.namespace, []);
+        }
+        return super.visitImportFrom(node);
+    }
 }
diff --git a/packages/pyright-internal/src/languageService/semanticTokensProvider.ts b/packages/pyright-internal/src/languageService/semanticTokensProvider.ts
index 71fa38ed1ac8..509d7e5c8b68 100644
--- a/packages/pyright-internal/src/languageService/semanticTokensProvider.ts
+++ b/packages/pyright-internal/src/languageService/semanticTokensProvider.ts
@@ -19,6 +19,7 @@ export const tokenTypes: string[] = [
     SemanticTokenTypes.method,
     SemanticTokenTypes.decorator,
     SemanticTokenTypes.property,
+    SemanticTokenTypes.namespace,
 ];
 export const tokenModifiers: string[] = [
     SemanticTokenModifiers.definition,