refactor: import action by Pyright's autoimport

fannheyward · Nov 17, 2023 · 61781af · 61781af
1 parent f7199c3
commit 61781af
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 91 deletions.
diff --git a/src/features/codeAction.ts b/src/features/codeAction.ts
@@ -1,4 +1,19 @@
-import { CodeAction, CodeActionContext, CodeActionKind, CodeActionProvider, Diagnostic, Position, ProviderResult, Range, TextDocument, TextEdit, workspace } from 'coc.nvim';
+import {
+  CancellationTokenSource,
+  CodeAction,
+  CodeActionContext,
+  CodeActionKind,
+  CodeActionProvider,
+  CompleteOption,
+  Diagnostic,
+  Position,
+  Range,
+  TextDocument,
+  TextEdit,
+  VimCompleteItem,
+  sources,
+  workspace,
+} from 'coc.nvim';
 
 export class PythonCodeActionProvider implements CodeActionProvider {
   private wholeRange(doc: TextDocument, range: Range): boolean {
@@ -95,55 +110,64 @@ export class PythonCodeActionProvider implements CodeActionProvider {
     ];
   }
 
-  private fixAction(document: TextDocument, diag: Diagnostic): CodeAction[] {
-    if (diag.code === 'reportUndefinedVariable') {
-      const msg = diag.message;
-      const match = msg.match(/"(.*)" is not defined/);
-      if (match) {
-        return [
-          {
-            title: `Add "import ${match[1]}"`,
-            kind: CodeActionKind.QuickFix,
-            command: {
-              title: '',
-              command: 'pyright.addImport',
-              arguments: [document, match[1], false],
-            },
-          },
+  private async fetchImportsByDiagnostic(document: TextDocument, diag: Diagnostic): Promise<ReadonlyArray<VimCompleteItem>> {
+    const match = diag.message.match(/"(.*)" is not defined/);
+    if (!match) return [];
+
+    const source = sources.sources.find((s) => s.name.includes('pyright'));
+    if (!source) return [];
+
+    // @ts-ignore
+    const option: CompleteOption = { position: diag.range.end, bufnr: document.uri };
+    const tokenSource = new CancellationTokenSource();
+    const result = await source.doComplete(option, tokenSource.token);
+    tokenSource.cancel();
+
+    // @ts-ignore
+    return result ? result.items.filter(x => x.label === match[1]) : [];
+  }
 
-          {
-            title: `Add "from _ import ${match[1]}"`,
-            kind: CodeActionKind.QuickFix,
-            command: {
-              title: '',
-              command: 'pyright.addImport',
-              arguments: [document, match[1], true],
+  private async fixAction(document: TextDocument, diag: Diagnostic): Promise<CodeAction[]> {
+    const actions: CodeAction[] = [];
+    if (diag.code === 'reportUndefinedVariable') {
+      const items = await this.fetchImportsByDiagnostic(document, diag);
+      for (const item of items) {
+        // @ts-ignore
+        const changes: TextEdit[] = [item.textEdit].concat(item.additionalTextEdits ?? []);
+        // @ts-ignore
+        const title = item.documentation?.value.replace('```\n', '').replace('\n```', '').trim();
+        actions.push({
+          title,
+          kind: CodeActionKind.QuickFix,
+          edit: {
+            changes: {
+              [document.uri]: changes,
             },
           },
-        ];
+        });
       }
     }
 
     // @ts-ignore
     if (diag.fix) {
-      const action: CodeAction = {
+      actions.push({
         // @ts-ignore
         title: diag.fix.title,
         kind: CodeActionKind.QuickFix,
         // @ts-ignore
         edit: diag.fix.edit,
-      };
-      return [action];
+      });
     }
-    return [];
+
+    return actions;
   }
 
-  public provideCodeActions(document: TextDocument, range: Range, context: CodeActionContext): ProviderResult<CodeAction[]> {
+  public async provideCodeActions(document: TextDocument, range: Range, context: CodeActionContext): Promise<CodeAction[] | null | undefined> {
     const actions: CodeAction[] = [];
 
     if (context.diagnostics.length) {
       for (const diag of context.diagnostics) {
-        actions.push(...this.fixAction(document, diag));
+        actions.push(...await this.fixAction(document, diag));
       }
     }
 

diff --git a/src/features/importCompletion.ts b/src/features/importCompletion.ts
@@ -14,7 +14,6 @@ export class ImportCompletionProvider implements CompletionItemProvider {
     }
     const source = sources.sources.find((s) => s.name.includes('pyright'));
     if (!source) return [];
-    // @ts-ignore
     const result = await source.doComplete(context.option, token);
     if (!result) return [];
     const items: CompletionItem[] = [];

diff --git a/src/features/refactor.ts b/src/features/refactor.ts
@@ -255,55 +255,6 @@ export async function extractMethod(root: string, document: TextDocument, range:
   });
 }
 
-export async function addImport(root: string, document: TextDocument, name: string, parent: boolean, outputChannel: OutputChannel): Promise<void> {
-  const pythonToolsExecutionService = new PythonExecutionService();
-  const rope = await pythonToolsExecutionService.isModuleInstalled('rope');
-  if (!rope) {
-    window.showWarningMessage(`Module rope not installed`);
-    return;
-  }
-
-  const doc = workspace.getDocument(document.uri);
-  const tempFile = await getTempFileWithDocumentContents(document);
-  let parentModule = '';
-  if (parent) parentModule = await window.requestInput('Module:');
-
-  const workspaceFolder = workspace.getWorkspaceFolder(doc.uri);
-  const workspaceRoot = workspaceFolder ? Uri.parse(workspaceFolder.uri).fsPath : workspace.cwd;
-  const pythonSettings = PythonSettings.getInstance();
-  return validateDocumentForRefactor(doc).then(() => {
-    const proxy = new RefactorProxy(root, pythonSettings, workspaceRoot);
-    const resp = proxy.addImport<RenameResponse>(doc.textDocument, tempFile, name, parentModule).then((response) => {
-      return response.results[0].diff;
-    });
-    return applyImports(doc, resp, outputChannel, tempFile);
-  });
-}
-
-async function applyImports(doc: Document, resp: Promise<string>, outputChannel: OutputChannel, tempFile: string): Promise<any> {
-  try {
-    const diff = await resp;
-    if (diff.length === 0) return;
-
-    const edits = getTextEditsFromPatch(doc.getDocumentContent(), diff);
-    await doc.applyEdits(edits);
-    await fs.promises.unlink(tempFile);
-  } catch (error) {
-    let errorMessage = `${error}`;
-    if (typeof error === 'string') {
-      errorMessage = error;
-    }
-    if (error instanceof Error) {
-      errorMessage = error.message;
-    }
-    outputChannel.appendLine(`${'#'.repeat(10)}Rope Output${'#'.repeat(10)}`);
-    outputChannel.appendLine(`Error in add import:\n${errorMessage}`);
-    outputChannel.appendLine('');
-    window.showErrorMessage(`Cannot perform addImport using selected element(s).`);
-    return await Promise.reject(error);
-  }
-}
-
 async function extractName(textEditor: Document, newName: string, renameResponse: Promise<string>, outputChannel: OutputChannel, tempFile: string): Promise<any> {
   let changeStartsAtLine = -1;
   try {

diff --git a/src/index.ts b/src/index.ts
@@ -29,7 +29,7 @@ import { TypeInlayHintsProvider } from './features/inlayHints';
 import { PythonSemanticTokensProvider } from './features/semanticTokens';
 import { sortImports } from './features/sortImports';
 import { LinterProvider } from './features/lintting';
-import { addImport, extractMethod, extractVariable } from './features/refactor';
+import { extractMethod, extractVariable } from './features/refactor';
 import { TestFrameworkProvider } from './features/testing';
 import { configuration, handleDiagnostics, provideCompletionItem, provideHover, provideSignatureHelp, resolveCompletionItem } from './middleware';
 
@@ -222,16 +222,6 @@ export async function activate(context: ExtensionContext): Promise<void> {
   );
   context.subscriptions.push(disposable);
 
-  disposable = commands.registerCommand(
-    'pyright.addImport',
-    async (document: TextDocument, name: string, parent: boolean) => {
-      await addImport(context.extensionPath, document, name, parent, outputChannel).catch(() => {});
-    },
-    null,
-    true
-  );
-  context.subscriptions.push(disposable);
-
   disposable = commands.registerCommand('python.sortImports', async () => {
     await sortImports(outputChannel).catch(() => {});
   });