From a9469b0ccaf59807e7fb03e07e471e137208b97a Mon Sep 17 00:00:00 2001 From: detachhead Date: Wed, 25 Dec 2024 18:49:31 +1000 Subject: [PATCH] wip --- .../src/analyzer/typeInlayHintsWalker.ts | 84 ++++++++++++++++--- .../src/analyzer/typePrinter.ts | 19 +++++ .../src/languageService/autoImporter.ts | 7 +- 3 files changed, 95 insertions(+), 15 deletions(-) diff --git a/packages/pyright-internal/src/analyzer/typeInlayHintsWalker.ts b/packages/pyright-internal/src/analyzer/typeInlayHintsWalker.ts index 9680301bfb..aadcea5017 100644 --- 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 index 5f4da5f1bf..f02641ca13 100644 --- 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 index be034ec93e..44f5f6da63 100644 --- 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,