wip

DetachHead · Dec 25, 2024 · a9469b0 · a9469b0
1 parent aba927d
commit a9469b0
Show file tree

Hide file tree

Showing 3 changed files with 95 additions and 15 deletions.
diff --git a/packages/pyright-internal/src/analyzer/typeInlayHintsWalker.ts b/packages/pyright-internal/src/analyzer/typeInlayHintsWalker.ts
@@ -1,10 +1,13 @@
-import { Range } from 'vscode-languageserver-types';
+import { Range, TextEdit } from 'vscode-languageserver-types';
 import { ParseTreeWalker } from '../analyzer/parseTreeWalker';
 import { isDunderName, isUnderscoreOnlyName } from '../analyzer/symbolNameUtils';
 import {
+    AnyType,
     ClassType,
     FunctionType,
+    NeverType,
     Type,
+    UnknownType,
     getTypeAliasInfo,
     isAny,
     isClass,
@@ -33,11 +36,20 @@ import { Uri } from '../common/uri/uri';
 import { ParseFileResults } from '../parser/parser';
 import { transformTypeForEnumMember } from './enums';
 import { InlayHintSettings } from '../workspaceFactory';
+import { AutoImporter } from '../languageService/autoImporter';
+import { printTypeAndGetRequiredImports, PrintTypeFlags } from './typePrinter';
+import { ImportGroup } from './importStatementUtils';
+import { TypeWalker } from './typeWalker';
+import { FileUri } from '../common/uri/fileUri';
+import { Program } from './program';
+
+type InlayHintType = 'variable' | 'functionReturn' | 'parameter' | 'generic';
 
 export type TypeInlayHintsItemType = {
-    inlayHintType: 'variable' | 'functionReturn' | 'parameter' | 'generic';
+    inlayHintType: InlayHintType;
     position: number;
     value: string;
+    imports: TextEdit[];
 };
 // Don't generate inlay hints for arguments to builtin types and functions
 const ignoredBuiltinTypes = new Set(
@@ -94,6 +106,42 @@ function isLeftSideOfAssignment(node: ParseNode): boolean {
     return node.start < node.parent.d.rightExpr.start;
 }
 
+/**
+ * tracks what imports need to be added when an inlay hint gets converted to real life by double clicking it
+ */
+class InlayHintImportFinder extends TypeWalker {
+    private _typingModule: Uri;
+    imports = new Array<{ name: string; from: Uri }>();
+
+    constructor(program: Program) {
+        super();
+        const userFiles = program.getUserFiles();
+        this._typingModule = userFiles.find((file) => file.sourceFile.isTypingStubFile())!.sourceFile.getUri();
+    }
+
+    override visitAny(type: AnyType): void {
+        this._addAny();
+        super.visitAny(type);
+    }
+    override visitUnknown(type: UnknownType): void {
+        // Unknown isn't a real type so the inlay hints show Any instead
+        this._addAny();
+        super.visitUnknown(type);
+    }
+    override visitNever(type: NeverType): void {
+        this.imports.push({ name: type.priv.isNoReturn ? 'NoReturn' : 'Never', from: this._typingModule });
+        super.visitNever(type);
+    }
+
+    override visitFunction(type: FunctionType): void {
+        // TODO: are all instances of FunctionType represented as `Callable`?
+        this.imports.push({ name: 'Callable', from: this._typingModule });
+        super.visitFunction(type);
+    }
+
+    private _addAny = () => this.imports.push({ name: 'Any', from: this._typingModule });
+}
+
 export class TypeInlayHintsWalker extends ParseTreeWalker {
     featureItems: TypeInlayHintsItemType[] = [];
     parseResults?: ParseFileResults;
@@ -103,6 +151,7 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
     constructor(
         private readonly _program: ProgramView,
         private _settings: InlayHintSettings,
+        private _autoImporter: AutoImporter,
         fileUri: Uri,
         range?: Range
     ) {
@@ -153,17 +202,18 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
                 !isTypeVar(type) &&
                 !isParamSpec(type)
             ) {
-                this.featureItems.push({
-                    inlayHintType: 'variable',
-                    position: this._endOfNode(node),
-                    value: `: ${
-                        type.props?.typeAliasInfo &&
+                const { printedType, imports } = this._printType(
+                    type.props?.typeAliasInfo &&
                         node.nodeType === ParseNodeType.Name &&
                         // prevent variables whose type comes from a type alias from being incorrectly treated as a TypeAlias.
                         getTypeAliasInfo(type)?.shared.name === node.d.value
-                            ? 'TypeAlias'
-                            : this._printType(type)
-                    }`,
+                        ? this._program.evaluator!.getBuiltInType(node, 'TypeAlias')
+                        : type
+                );
+                this.featureItems.push({
+                    inlayHintType: 'variable',
+                    position: this._endOfNode(node),
+                    value: `: ${printedType.printedType}`,
                 });
             }
         }
@@ -330,6 +380,16 @@ export class TypeInlayHintsWalker extends ParseTreeWalker {
 
     private _endOfNode = (node: ParseNode) => node.start + node.length;
 
-    private _printType = (type: Type): string =>
-        this._program.evaluator!.printType(type, { enforcePythonSyntax: true });
+    private _printType = (type: Type) =>
+        printTypeAndGetRequiredImports(type, PrintTypeFlags.PythonSyntax, (type) =>
+            this._program.evaluator!.getEffectiveReturnType(type)
+        );
+
+    private _importsToTextEdits = (imports: string[]) => {
+        for (const fullyQualifiedName of imports) {
+            const [_, name, module] = fullyQualifiedName.match(/(.*?)\.(.*)/)!;
+            // TODO: figure out the correct import group
+            this._autoImporter.getTextEditsForAutoImportByFilePath({ name }, { module }, name, ImportGroup.BuiltIn);
+        }
+    };
 }
diff --git a/packages/pyright-internal/src/analyzer/typePrinter.ts b/packages/pyright-internal/src/analyzer/typePrinter.ts
@@ -33,6 +33,7 @@ import {
     OverloadedType,
     TupleTypeArg,
     Type,
+    TypeAliasInfo,
     TypeBase,
     TypeCategory,
     TypeVarType,
@@ -93,6 +94,23 @@ export const enum PrintTypeFlags {
 
 export type FunctionReturnTypeCallback = (type: FunctionType) => Type;
 
+/**
+ * useful if the printed type neds to be inserted into the code via a `TextEdit`, eg. when inserting an inlay hint
+ * where it needs to also insert imports
+ */
+export const printTypeAndGetRequiredImports = (
+    type: Type,
+    printTypeFlags: PrintTypeFlags,
+    returnTypeCallback: FunctionReturnTypeCallback
+) => {
+    const uniqueNameMap = new UniqueNameMap(printTypeFlags | PrintTypeFlags.UseFullyQualifiedNames, returnTypeCallback);
+    uniqueNameMap.build(type);
+    return {
+        printedType: printType(type, printTypeFlags, returnTypeCallback),
+        imports: uniqueNameMap.map.keys(),
+    };
+};
+
 export function printType(
     type: Type,
     printTypeFlags: PrintTypeFlags,
@@ -183,6 +201,7 @@ export function printLiteralValue(type: ClassType, quotation = "'"): string {
     return literalStr;
 }
 
+// make this return an import tracker obkject or somethinmg......................
 function printTypeInternal(
     type: Type,
     printTypeFlags: PrintTypeFlags,

diff --git a/packages/pyright-internal/src/languageService/autoImporter.ts b/packages/pyright-internal/src/languageService/autoImporter.ts
@@ -306,7 +306,7 @@ export class AutoImporter {
                     return;
                 }
 
-                const autoImportTextEdits = this._getTextEditsForAutoImportByFilePath(
+                const autoImportTextEdits = this.getTextEditsForAutoImportByFilePath(
                     { name: importAliasData.importParts.symbolName, alias: abbrFromUsers },
                     {
                         name: importAliasData.importParts.importFrom ?? importAliasData.importParts.importName,
@@ -396,7 +396,7 @@ export class AutoImporter {
             }
 
             const nameForImportFrom = this.getNameForImportFrom(/* library */ !fileProperties.isUserCode, moduleUri);
-            const autoImportTextEdits = this._getTextEditsForAutoImportByFilePath(
+            const autoImportTextEdits = this.getTextEditsForAutoImportByFilePath(
                 { name, alias: abbrFromUsers },
                 { name: importSource, nameForImportFrom },
                 name,
@@ -638,7 +638,8 @@ export class AutoImporter {
         return this.importResolver.getModuleNameForImport(uri, this.execEnvironment);
     }
 
-    private _getTextEditsForAutoImportByFilePath(
+    // eslint-disable-next-line @typescript-eslint/member-ordering -- this is private upstream and to minimize conflicts im not moving it
+    getTextEditsForAutoImportByFilePath(
         importNameInfo: ImportNameInfo,
         moduleNameInfo: ModuleNameInfo,
         insertionText: string,