diff --git a/packages/pyright-internal/src/languageService/codeActionProvider.ts b/packages/pyright-internal/src/languageService/codeActionProvider.ts index 2a389ce9d..0870b7228 100644 --- a/packages/pyright-internal/src/languageService/codeActionProvider.ts +++ b/packages/pyright-internal/src/languageService/codeActionProvider.ts @@ -12,12 +12,12 @@ import { Commands } from '../commands/commands'; import { throwIfCancellationRequested } from '../common/cancellationUtils'; import { createCommand } from '../common/commandUtils'; import { CreateTypeStubFileAction } from '../common/diagnostic'; -import { Range } from '../common/textRange'; +import { Range, TextRange } from '../common/textRange'; import { Uri } from '../common/uri/uri'; import { convertToFileTextEdits, convertToTextEditActions, convertToWorkspaceEdit } from '../common/workspaceEditUtils'; import { Localizer } from '../localization/localize'; import { Workspace } from '../workspaceFactory'; -import { CompletionProvider } from './completionProvider'; +import { CompletionMap, CompletionProvider } from './completionProvider'; import { convertOffsetToPosition, convertPositionToOffset, @@ -25,9 +25,11 @@ import { getLineEndPosition, } from '../common/positionUtils'; import { findNodeByOffset } from '../analyzer/parseTreeUtils'; -import { ParseNodeType } from '../parser/parseNodes'; +import { FunctionNode, ParseNode, ParseNodeType } from '../parser/parseNodes'; import { sorter } from '../common/collectionUtils'; import { LanguageServerInterface } from '../common/languageServerInterface'; +import { DiagnosticRule } from '../common/diagnosticRules'; +import { TextRangeCollection } from '../common/textRangeCollection'; export class CodeActionProvider { static mightSupport(kinds: CodeActionKind[] | undefined): boolean { @@ -64,117 +66,139 @@ export class CodeActionProvider { const parseResults = workspace.service.backgroundAnalysisProgram.program.getParseResults(fileUri)!; const lines = parseResults.tokenizerOutput.lines; - if (diags.find((d) => d.getActions()?.some((action) => action.action === Commands.import))) { - const offset = convertPositionToOffset(range.start, lines); - if (offset === undefined) { - return []; + const fs = ls.serviceProvider.fs(); + for (const diagnostic of diags) { + const rule = diagnostic.getRule(); + if (!rule) { + continue; } + const line = diagnostic.range.start.line; - const node = findNodeByOffset(parseResults.parserOutput.parseTree, offset); - if (node === undefined) { - return []; - } + // ==== CA for auto imports ==== + if (diagnostic.getActions()?.some((action) => action.action === Commands.import)) { + const offset = convertPositionToOffset(range.start, lines); + if (offset === undefined) { + return []; + } - const completer = new CompletionProvider( - workspace.service.backgroundAnalysisProgram.program, - fileUri, - convertOffsetToPosition(node.start + node.length, lines), - { - format: 'plaintext', - lazyEdit: false, - snippet: false, - // we don't care about deprecations here so make sure they don't get evaluated unnecessarily - // (we don't call resolveCompletionItem) - checkDeprecatedWhenResolving: true, - useTypingExtensions: workspace.useTypingExtensions, - }, - token, - true - ); + const node = findNodeByOffset(parseResults.parserOutput.parseTree, offset); + if (node === undefined) { + return []; + } - const word = node.nodeType === ParseNodeType.Name ? node.d.value : undefined; - const sortedCompletions = completer - .getCompletions() - ?.items.filter( - (completion) => - // only show exact matches as code actions, which matches pylance's behavior. otherwise it's too noisy - // because code actions don't get sorted like completions do. see https://github.com/DetachHead/basedpyright/issues/747 - completion.label === word - ) - .sort((prev, next) => - sorter( - prev, - next, - (prev, next) => (prev.sortText && next.sortText && prev.sortText < next.sortText) || false + const completer = this._createCompleter(workspace, fileUri, token, node, lines); + + const word = node.nodeType === ParseNodeType.Name ? node.d.value : undefined; + const sortedCompletions = completer + .getCompletions() + ?.items.filter( + (completion) => + // only show exact matches as code actions, which matches pylance's behavior. otherwise it's too noisy + // because code actions don't get sorted like completions do. see https://github.com/DetachHead/basedpyright/issues/747 + completion.label === word ) - ); + .sort((prev, next) => + sorter( + prev, + next, + (prev, next) => (prev.sortText && next.sortText && prev.sortText < next.sortText) || false + ) + ); - for (const suggestedImport of sortedCompletions ?? []) { - if (!suggestedImport.data) { - continue; - } - let textEdits: TextEdit[] = []; - if (suggestedImport.textEdit && 'range' in suggestedImport.textEdit) { - textEdits.push(suggestedImport.textEdit); + for (const suggestedImport of sortedCompletions ?? []) { + if (!suggestedImport.data) { + continue; + } + let textEdits: TextEdit[] = []; + if (suggestedImport.textEdit && 'range' in suggestedImport.textEdit) { + textEdits.push(suggestedImport.textEdit); + } + if (suggestedImport.additionalTextEdits) { + textEdits = textEdits.concat(suggestedImport.additionalTextEdits); + } + if (textEdits.length === 0) { + continue; + } + const workspaceEdit = convertToWorkspaceEdit( + ls.convertUriToLspUriString, + completer.importResolver.fileSystem, + convertToFileTextEdits(fileUri, convertToTextEditActions(textEdits)) + ); + codeActions.push( + CodeAction.create( + suggestedImport.data.autoImportText.trim(), + workspaceEdit, + CodeActionKind.QuickFix + ) + ); } - if (suggestedImport.additionalTextEdits) { - textEdits = textEdits.concat(suggestedImport.additionalTextEdits); + } + + // ==== CA for adding @override decorators ==== + if (rule === DiagnosticRule.reportImplicitOverride) { + const lineText = lines.getItemAt(line); + const methodLineContent = parseResults.text.substring(lineText.start, lineText.start + lineText.length); + + const functionLine = { line: line, character: 0 }; + const offset = convertPositionToOffset(functionLine, lines); + if (offset === undefined) { + return []; } - if (textEdits.length === 0) { - continue; + + const node = findNodeByOffset(parseResults.parserOutput.parseTree, offset); + if (node === undefined) { + return []; } - const workspaceEdit = convertToWorkspaceEdit( - ls.convertUriToLspUriString, - completer.importResolver.fileSystem, - convertToFileTextEdits(fileUri, convertToTextEditActions(textEdits)) - ); - codeActions.push( - CodeAction.create( - suggestedImport.data.autoImportText.trim(), - workspaceEdit, - CodeActionKind.QuickFix - ) - ); - } - } - if (!workspace.rootUri) { - return codeActions; - } + const completer = this._createCompleter(workspace, fileUri, token, node, lines); + const completionMap = new CompletionMap(); - const typeStubDiag = diags.find((d) => { - const actions = d.getActions(); - return actions && actions.find((a) => a.action === Commands.createTypeStub); - }); - - if (typeStubDiag) { - const action = typeStubDiag - .getActions()! - .find((a) => a.action === Commands.createTypeStub) as CreateTypeStubFileAction; - if (action) { - const createTypeStubAction = CodeAction.create( - Localizer.CodeAction.createTypeStubFor().format({ moduleName: action.moduleName }), - createCommand( - Localizer.CodeAction.createTypeStub(), - Commands.createTypeStub, - workspace.rootUri.toString(), - action.moduleName, - fileUri.toString() - ), - CodeActionKind.QuickFix + const overrideEdits = completer.createOverrideEdits( + node as FunctionNode, + offset, + undefined, + methodLineContent, + completionMap ); - codeActions.push(createTypeStubAction); + + if (overrideEdits.length > 0) { + codeActions.push( + CodeAction.create( + Localizer.CodeAction.addExplicitOverride(), + convertToWorkspaceEdit( + ls.convertUriToLspUriString, + fs, + convertToFileTextEdits(fileUri, overrideEdits) + ), + CodeActionKind.QuickFix + ) + ); + } } - } - const fs = ls.serviceProvider.fs(); - for (const diagnostic of diags) { - const rule = diagnostic.getRule(); - if (!rule) { - continue; + // ==== CA for creating type stubs ==== + if (workspace.rootUri && diagnostic.getActions()?.some((a) => a.action === Commands.createTypeStub)) { + const action = diagnostic + .getActions()! + .find((a) => a.action === Commands.createTypeStub) as CreateTypeStubFileAction; + if (action) { + const createTypeStubAction = CodeAction.create( + Localizer.CodeAction.createTypeStubFor().format({ moduleName: action.moduleName }), + createCommand( + Localizer.CodeAction.createTypeStub(), + Commands.createTypeStub, + workspace.rootUri.toString(), + action.moduleName, + fileUri.toString() + ), + CodeActionKind.QuickFix + ); + codeActions.push(createTypeStubAction); + } } + + // ==== CA for adding ignore comments ==== const ignoreCommentPrefix = `# pyright: ignore`; - const line = diagnostic.range.start.line; // we deliberately only check for pyright:ignore comments here but not type:ignore for 2 reasons: // - type:ignore comments are discouraged in favor of pyright:ignore comments // - the type:ignore comment might be for another type checker @@ -189,12 +213,15 @@ export class CodeActionProvider { } positionCharacter = convertTextRangeToRange(lastRuleTextRange, lines).end.character; insertText = `, ${rule}`; - title = `Add \`${rule}\` to existing \`${ignoreCommentPrefix}\` comment`; + title = Localizer.CodeAction.addIgnoreCommentToExisting().format({ + rule: rule, + ignoreCommentPrefix: ignoreCommentPrefix, + }); } else { positionCharacter = getLineEndPosition(parseResults.tokenizerOutput, parseResults.text, line).character; const ignoreComment = `${ignoreCommentPrefix}[${rule}]`; insertText = ` ${ignoreComment}`; - title = `Add \`${ignoreComment}\``; + title = Localizer.CodeAction.addIgnoreComment().format({ ignoreComment: ignoreComment }); } const position = { line, character: positionCharacter }; codeActions.push( @@ -217,4 +244,29 @@ export class CodeActionProvider { return codeActions; } + + private static _createCompleter( + workspace: Workspace, + fileUri: Uri, + token: CancellationToken, + node: ParseNode, + lines: TextRangeCollection + ) { + return new CompletionProvider( + workspace.service.backgroundAnalysisProgram.program, + fileUri, + convertOffsetToPosition(node.start + node.length, lines), + { + format: 'plaintext', + lazyEdit: false, + snippet: false, + // we don't care about deprecations here so make sure they don't get evaluated unnecessarily + // (we don't call resolveCompletionItem) + checkDeprecatedWhenResolving: true, + useTypingExtensions: workspace.useTypingExtensions, + }, + token, + true + ); + } } diff --git a/packages/pyright-internal/src/languageService/completionProvider.ts b/packages/pyright-internal/src/languageService/completionProvider.ts index 502a7fb11..783a04cc7 100644 --- a/packages/pyright-internal/src/languageService/completionProvider.ts +++ b/packages/pyright-internal/src/languageService/completionProvider.ts @@ -100,6 +100,7 @@ import { ErrorNode, ExpressionNode, FormatStringNode, + FunctionNode, ImportFromNode, IndexNode, isExpressionNode, @@ -412,6 +413,81 @@ export class CompletionProvider { } } + createAutoImporter(completionMap: CompletionMap, lazyEdit: boolean) { + const currentFile = this.program.getSourceFileInfo(this.fileUri); + const moduleSymbolMap = buildModuleSymbolsMap( + this.program, + this.program.getSourceFileInfoList().filter((s) => s !== currentFile) + ); + + return new AutoImporter( + this.program, + this.execEnv, + this.parseResults, + this.position, + completionMap, + moduleSymbolMap, + { + lazyEdit, + } + ); + } + + createOverrideEdits( + node: FunctionNode, + startNode: number, + decorators: DecoratorNode[] | undefined, + priorText: string, + completionMap: CompletionMap + ) { + const additionalTextEdits: TextEditAction[] = []; + const overrideDecorator = this.evaluator.getTypingType(node, 'override'); + if ( + // this should always be true, but just in case + overrideDecorator?.category === TypeCategory.Function && + // if targeting a python version that doesn't have typing.override, we don't want to insert the decorator unless + // the user has explicitly enabled `basedpyright.analysis.useTypingExtensions` + (overrideDecorator.shared.moduleName !== 'typing_extensions' || this.options.useTypingExtensions) && + // check if the override decorator is already here + !decorators?.some((decorator) => { + const type = this.evaluator.getTypeOfExpression(decorator.d.expr).type; + return isFunction(type) && FunctionType.isBuiltIn(type, 'override'); + }) + ) { + const indent = priorText.match(/(^\s*)(async|def)\s/)?.[1]; + // this should always match, but it looks gross and hacky so just in case, + // we just skip adding the override decorator if it can't figure out where to put it + if (indent !== undefined) { + const position: Position = { + line: + // if there are any other decorators, make this the last one because some decorators can change the type + // such that the override decorator won't can't be placed above them + convertOffsetToPosition(startNode, this.parseResults.tokenizerOutput.lines).line + + (decorators?.length ?? 0), + character: indent.length, + }; + const importTextEditInfo = this.createAutoImporter( + completionMap, + this.options.lazyEdit + ).getTextEditsForAutoImportByFilePath( + { name: 'override' }, + { name: overrideDecorator.shared.moduleName }, + 'override', + ImportGroup.BuiltIn, + overrideDecorator.shared.declaration!.uri + ); + if (importTextEditInfo.edits) { + additionalTextEdits.push(...importTextEditInfo.edits); + } + additionalTextEdits.push({ + range: { start: position, end: position }, + replacementText: `@${importTextEditInfo.insertionText}\n${indent}`, + }); + } + } + return additionalTextEdits; + } + protected get evaluator() { return this.program.evaluator!; } @@ -529,52 +605,15 @@ export class CompletionProvider { // If reportImplicitOverride is disabled, never add @override this.configOptions.diagnosticRuleSet.reportImplicitOverride !== 'none' ) { - const overrideDecorator = this.evaluator.getTypingType(decl.node, 'override'); - if ( - // this should always be true, but just in case - overrideDecorator?.category === TypeCategory.Function && - // if targeting a python version that doesn't have typing.override, we don't want to insert the decorator unless - // the user has explicitly enabled `basedpyright.analysis.useTypingExtensions` - (overrideDecorator.shared.moduleName !== 'typing_extensions' || - this.options.useTypingExtensions) && - // check if the override decorator is already here - !decorators?.some((decorator) => { - const type = this.evaluator.getTypeOfExpression(decorator.d.expr).type; - return isFunction(type) && FunctionType.isBuiltIn(type, 'override'); - }) - ) { - const indent = priorText.match(/(^\s*)(async|def)\s/)?.[1]; - // this should always match, but it looks gross and hacky so just in case, - // we just skip adding the override decorator if it can't figure out where to put it - if (indent !== undefined) { - const startNode = partialName.parent!.start; - const position: Position = { - line: - // if there are any other decorators, make this the last one because some decorators can change the type - // such that the override decorator won't can't be placed above them - convertOffsetToPosition(startNode, this.parseResults.tokenizerOutput.lines) - .line + (decorators?.length ?? 0), - character: indent.length, - }; - const importTextEditInfo = this.createAutoImporter( - completionMap, - this.options.lazyEdit - ).getTextEditsForAutoImportByFilePath( - { name: 'override' }, - { name: overrideDecorator.shared.moduleName }, - 'override', - ImportGroup.BuiltIn, - overrideDecorator.shared.declaration!.uri - ); - if (importTextEditInfo.edits) { - additionalTextEdits.push(...importTextEditInfo.edits); - } - additionalTextEdits.push({ - range: { start: position, end: position }, - replacementText: `@${importTextEditInfo.insertionText}\n${indent}`, - }); - } - } + additionalTextEdits.push( + ...this.createOverrideEdits( + decl.node, + partialName.parent!.start, + decorators, + priorText, + completionMap + ) + ); } this.addSymbol(name, symbol, partialName.d.value, completionMap, { @@ -922,26 +961,6 @@ export class CompletionProvider { return completionMap; } - protected createAutoImporter(completionMap: CompletionMap, lazyEdit: boolean) { - const currentFile = this.program.getSourceFileInfo(this.fileUri); - const moduleSymbolMap = buildModuleSymbolsMap( - this.program, - this.program.getSourceFileInfoList().filter((s) => s !== currentFile) - ); - - return new AutoImporter( - this.program, - this.execEnv, - this.parseResults, - this.position, - completionMap, - moduleSymbolMap, - { - lazyEdit, - } - ); - } - protected addAutoImportCompletions( priorWord: string, similarityLimit: number, diff --git a/packages/pyright-internal/src/localization/localize.ts b/packages/pyright-internal/src/localization/localize.ts index 73d53a202..6e0dad15f 100644 --- a/packages/pyright-internal/src/localization/localize.ts +++ b/packages/pyright-internal/src/localization/localize.ts @@ -1705,6 +1705,13 @@ export namespace Localizer { export const findingReferences = () => getRawString('CodeAction.findingReferences'); export const findingImplementations = () => getRawString('CodeAction.findingImplementations'); export const organizeImports = () => getRawString('CodeAction.organizeImports'); + export const addExplicitOverride = () => getRawString('CodeAction.addExplicitOverride'); + export const addIgnoreComment = () => + new ParameterizedString<{ ignoreComment: string }>(getRawString('CodeAction.addIgnoreComment')); + export const addIgnoreCommentToExisting = () => + new ParameterizedString<{ rule: string; ignoreCommentPrefix: string }>( + getRawString('CodeAction.addIgnoreCommentToExisting') + ); } export namespace Completion { diff --git a/packages/pyright-internal/src/localization/package.nls.en-us.json b/packages/pyright-internal/src/localization/package.nls.en-us.json index 9f58dc376..1169045ae 100644 --- a/packages/pyright-internal/src/localization/package.nls.en-us.json +++ b/packages/pyright-internal/src/localization/package.nls.en-us.json @@ -13,7 +13,10 @@ "filesToAnalyzeOne": "1 file to analyze", "findingReferences": "Finding references", "findingImplementations": "Finding implementations", - "organizeImports": "Organize Imports" + "organizeImports": "Organize Imports", + "addExplicitOverride": "Add `@override`", + "addIgnoreComment": "Add `{ignoreComment}`", + "addIgnoreCommentToExisting": "Add `{rule}` to existing `{ignoreCommentPrefix}` comment" }, "Completion": { "autoImportDetail": "Auto-import", diff --git a/packages/pyright-internal/src/tests/fourslash/addOverride.invokeCodeAction.fourslash.ts b/packages/pyright-internal/src/tests/fourslash/addOverride.invokeCodeAction.fourslash.ts new file mode 100644 index 000000000..d3531af19 --- /dev/null +++ b/packages/pyright-internal/src/tests/fourslash/addOverride.invokeCodeAction.fourslash.ts @@ -0,0 +1,46 @@ +/// + +// @filename: pyrightconfig.json +//// { +//// "autoImportCompletions": false +//// "typeCheckingMode": "recommended" +//// } + +// @filename: test.py +//// [|/*import*/|]class Base: +//// def method(self) -> None: ... +//// +//// class Derived(Base): +//// [|/*override*/|]def method(self) -> None: ... + +{ + const importRange = helper.getPositionRange('import'); + const overrideRange = helper.getPositionRange('override'); + + const codeActionTitle = 'Add `@override`'; + + const codeActions = { + codeActions: [ + { + title: codeActionTitle, + edit: { + changes: { + 'file:///test.py': [ + // Assumes method is indented with 4 spaces. Adjust if necessary. + { range: overrideRange, newText: '@override\n ' }, + { range: importRange, newText: 'from typing import override\n\n\n' }, + ], + }, + }, + kind: 'quickfix', + }, + ], + }; + + //@ts-expect-error https://github.com/DetachHead/basedpyright/issues/86 + await helper.verifyCodeActions('included', { + noImportMarker: { codeActions: [] }, + startMarker: codeActions, + endMarker: codeActions, + }); +}