prevent inlay hints from appearing on enum members

Author: DetachHead

packages/pyright-internal/src/analyzer/typeInlayHintsWalker.ts

@@ -2,6 +2,7 @@ import { Range } from 'vscode-languageserver-types';
 import { ParseTreeWalker } from '../analyzer/parseTreeWalker';
 import { isDunderName, isUnderscoreOnlyName } from '../analyzer/symbolNameUtils';
 import {
+    ClassType,
     FunctionType,
     Type,
     getTypeAliasInfo,
@@ -15,6 +16,7 @@ import { ProgramView } from '../common/extensibility';
 import { limitOverloadBasedOnCall } from '../languageService/tooltipUtils';
 import {
     CallNode,
+    ClassNode,
     FunctionNode,
     NameNode,
     ParamCategory,
@@ -28,6 +30,7 @@ import { convertRangeToTextRange } from '../common/positionUtils';
 import { Uri } from '../common/uri/uri';
 import { ParseFileResults } from '../parser/parser';
 import { InlayHintSettings } from '../common/languageServerInterface';
+import { transformTypeForEnumMember } from './enums';
 
 export type TypeInlayHintsItemType = {
     inlayHintType: 'variable' | 'functionReturn' | 'parameter';
@@ -93,6 +96,7 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
     featureItems: TypeInlayHintsItemType[] = [];
     parseResults?: ParseFileResults;
     private _range: TextRange | undefined;
+    private _variablesThatShouldntHaveInlayHints = new Set<ParseNode>();
 
     constructor(
         private readonly _program: ProgramView,
@@ -110,6 +114,26 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
         }
     }
 
+    override visitClass(node: ClassNode): boolean {
+        const evaluator = this._program.evaluator;
+        if (evaluator) {
+            const classType = evaluator.getTypeOfClass(node)?.classType;
+            // prevent inlay hints from appearing on enum members
+            if (classType && ClassType.isEnumClass(classType)) {
+                ClassType.getSymbolTable(classType).forEach((symbol, name) => {
+                    const symbolType = transformTypeForEnumMember(evaluator, classType, name, true);
+                    if (symbolType) {
+                        const nameNode = symbol.getDeclarations()[0]?.node;
+                        if (nameNode) {
+                            this._variablesThatShouldntHaveInlayHints.add(nameNode);
+                        }
+                    }
+                });
+            }
+        }
+        return super.visitClass(node);
+    }
+
     override visitName(node: NameNode): boolean {
         if (
             this._settings.variableTypes &&
@@ -125,7 +149,8 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
                 !(isClass(type) && isLiteralType(type)) &&
                 !isTypeVar(type) &&
                 // !isFunction(type) &&
-                !isParamSpec(type)
+                !isParamSpec(type) &&
+                !this._variablesThatShouldntHaveInlayHints.has(node)
             ) {
                 this.featureItems.push({
                     inlayHintType: 'variable',

packages/pyright-internal/src/tests/samples/inlay_hints/variables.py

@@ -1,3 +1,4 @@
+from enum import Enum
 from typing import Literal
 from typing_extensions import TypeVar, ParamSpec
 
@@ -14,4 +15,7 @@ def asdf(a: Foo, b: type[Foo], c: int | str, d: Literal[1, 2], e: type[int]) ->
     bar = b # inlay hint (type, but not TypeAlias)
     baz = c # inlay hint
     qux = d # inlay hint
-    quxx = e # inlay hint (type, but not TypeAlias)
+    quxx = e # inlay hint (type, but not TypeAlias)
+
+class Foo(Enum):
+    a = 1 # no inlay hint

packages/pyright-internal/src/tests/typeInlayHintsWalker.test.ts

@@ -7,37 +7,37 @@ if (process.platform !== 'win32' || !process.env['CI']) {
         expect(result).toStrictEqual([
             {
                 inlayHintType: 'variable',
-                position: 80,
+                position: 102,
                 value: ': str',
             },
             {
                 inlayHintType: 'variable',
-                position: 386,
+                position: 408,
                 value: ': TypeAlias',
             },
             {
                 inlayHintType: 'variable',
-                position: 547,
+                position: 569,
                 value: ': Foo',
             },
             {
                 inlayHintType: 'variable',
-                position: 572,
+                position: 594,
                 value: ': type[int]',
             },
             {
                 inlayHintType: 'variable',
-                position: 623,
+                position: 645,
                 value: ': int | str',
             },
             {
                 inlayHintType: 'variable',
-                position: 648,
+                position: 670,
                 value: ': Literal[1, 2]',
             },
             {
                 inlayHintType: 'variable',
-                position: 674,
+                position: 696,
                 value: ': type[int]',
             },
         ]);