#764 feature request: auto import (DonJayamanne#1032)

* Auto-import command for a symbol under cursor. Using symbol provider to locate where the symbol can be imported from. To improve: Position of imported line Importing modules Relative imports * Added tests for autoimport command * Add sinon to dev-package dependencies sinon is very popular mock library (http://sinonjs.org/), also used in vscode tests.

Author: skyjur

package.json

@@ -1506,11 +1506,13 @@
 "@types/socket.io": "^1.4.27",
 "@types/socket.io-client": "^1.4.27",
 "@types/uuid": "^3.3.27",
+"@types/sinon": "^2.3.2",
 "babel-core": "^6.14.0",
 "babel-loader": "^6.2.5",
 "babel-preset-es2015": "^6.14.0",
 "ignore-loader": "^0.1.1",
 "mocha": "^2.3.3",
+"sinon": "^2.3.6",
 "retyped-diff-match-patch-tsd-ambient": "^1.0.0-0",
 "ts-loader": "^0.8.2",
 "tslint": "^3.15.1",

src/client/autoImport.ts

@@ -0,0 +1,12 @@
+'use strict';
+
+import * as vscode from 'vscode';
+import * as fs from 'fs';
+import * as path from 'path';
+import { AutoImportProvider } from './providers/autoImportProvider';
+
+
+export function activate(context: vscode.ExtensionContext, outChannel: vscode.OutputChannel) {
+ let extension = new AutoImportProvider();
+ context.subscriptions.push(vscode.commands.registerCommand('python.autoImportAtCursor', extension.autoImportAtCursor));
+}

src/client/extension.ts

@@ -8,6 +8,7 @@ import { PythonReferenceProvider } from './providers/referenceProvider';
 import { PythonRenameProvider } from './providers/renameProvider';
 import { PythonFormattingEditProvider } from './providers/formatProvider';
 import * as sortImports from './sortImports';
+import * as autoImport from './autoImport';
 import { LintProvider } from './providers/lintProvider';
 import { PythonSymbolProvider } from './providers/symbolProvider';
 import { PythonSignatureProvider } from './providers/signatureProvider';
@@ -58,6 +59,7 @@ export function activate(context: vscode.ExtensionContext) {
  }
 
  sortImports.activate(context, formatOutChannel);
+ autoImport.activate(context, formatOutChannel);
  context.subscriptions.push(activateSetInterpreterProvider());
  context.subscriptions.push(...activateExecInTerminalProvider());
  context.subscriptions.push(activateUpdateSparkLibraryProvider());

src/client/providers/autoImportProvider.ts

@@ -0,0 +1,85 @@
+import * as vscode from 'vscode';
+import * as fs from 'fs';
+import * as path from 'path';
+
+export class AutoImportProvider {
+ 
+ constructor() {
+ }
+
+ autoImportAtCursor() {
+ let ed = vscode.window.activeTextEditor;
+ let range : vscode.Range = ed.selection;
+
+ if(range.start.line !== range.end.line) {
+ return;
+ }
+ 
+ if(range.start.character === range.end.character) {
+ range = ed.document.getWordRangeAtPosition(range.end);
+ }
+ const symbol = vscode.window.activeTextEditor.document.getText(range);
+
+ if(!symbol) {
+ return;
+ }
+
+ return this.autoImport(symbol);
+ }
+
+ async autoImport(symbol: string) {
+ let editor = vscode.window.activeTextEditor;
+ let result = await this.lookupSymbol(symbol);
+ 
+ result = result.filter((s: vscode.SymbolInformation) =>
+ s.name === symbol && // Only exact and case sensitive matches should be considered for import
+ s.kind !== vscode.SymbolKind.Namespace // only declarations should be considered for import
+ );
+ 
+ if(result.length === 0) {
+ vscode.window.showInformationMessage('No matching symbols found');
+ return;
+ } else {
+ var import_choices: string[] = result.map(
+ s => `from ${pathAsPyModule(s.location)} import ${s.name}`
+ );
+
+ let s = await this.showChoices(import_choices);
+ if(s) {
+ return addImport(editor, s); 
+ }
+ }
+ }
+
+ lookupSymbol(symbol: string) {
+ return <Promise<vscode.SymbolInformation[]>>
+ vscode.commands.executeCommand('vscode.executeWorkspaceSymbolProvider', symbol);
+ }
+
+ showChoices(import_choices: string[]) {
+ return vscode.window.showQuickPick(import_choices);
+ }
+}
+
+export function pathAsPyModule(l: vscode.Location): string {
+ var pymodule = path.basename(l.uri.fsPath).replace(/\.py$/, '');
+ var location = path.dirname(l.uri.fsPath);
+ while(fs.existsSync(path.join(location, '__init__.py'))) {
+ pymodule = path.basename(location) + '.' + pymodule;
+ location = path.dirname(location);
+ }
+ return pymodule;
+}
+
+export function addImport(ed: vscode.TextEditor, import_string: string) {
+ return ed.edit((b: vscode.TextEditorEdit) => b.insert(
+ getPositionForNewImport(import_string),
+ import_string + '\n'
+ ));
+}
+
+
+export function getPositionForNewImport(import_string: string): vscode.Position {
+ // TODO: figure out better position:
+ return new vscode.Position(0, 0);
+}

src/test/extension.autoimport.test.ts

@@ -0,0 +1,81 @@
+import * as assert from 'assert';
+import * as sinon from 'sinon';
+import * as vscode from 'vscode';
+import * as path from 'path';
+import { AutoImportProvider } from './../client/providers/autoImportProvider';
+
+let testPath = path.join(__dirname, '../../src/test');
+const fileOne = path.join(testPath, 'pythonFiles/autoimport/one.py');
+const fileThree = path.join(testPath, 'pythonFiles/autoimport/two/three.py');
+
+let fileOneLoc = new vscode.Location(vscode.Uri.file(fileOne), new vscode.Position(0, 0));
+let fileThreeLoc = new vscode.Location(vscode.Uri.file(fileThree), new vscode.Position(0, 0));
+
+
+suite('Autoimport', () => {
+ let autoimp = new AutoImportProvider();
+ let lookupSymbol = sinon.stub(autoimp, 'lookupSymbol');
+ let showChoices = sinon.stub(autoimp, 'showChoices');
+
+ suiteSetup(() => {
+ });
+
+ setup(() => {
+ lookupSymbol.reset();
+ showChoices.reset();
+ return vscode.workspace.openTextDocument(fileOne).then(vscode.window.showTextDocument); 
+ });
+
+ test('choices are filtered and displayed"', done => {
+ lookupSymbol.resolves([
+ new vscode.SymbolInformation(
+ 'TestClass', vscode.SymbolKind.Class, '', fileOneLoc
+ ),
+ new vscode.SymbolInformation(
+ 'TestClassTwo', vscode.SymbolKind.Class, '', fileOneLoc
+ ),
+ new vscode.SymbolInformation(
+ 'TestClass', vscode.SymbolKind.Namespace, '', fileOneLoc
+ ),
+ ]);
+ 
+ showChoices.resolves(undefined); 
+
+ autoimp.autoImport('TestClass').then(function() {
+ assert(lookupSymbol.calledWith('TestClass'));
+ assert(showChoices.calledWith(['from one import TestClass']));
+ done();
+ });
+ });
+
+ test('module under a package"', done => {
+ lookupSymbol.resolves([
+ new vscode.SymbolInformation(
+ 'TestClass', vscode.SymbolKind.Class, '', fileThreeLoc
+ ),
+ ]);
+
+ showChoices.resolves(undefined);
+
+ autoimp.autoImport('TestClass').then(function() {
+ assert(showChoices.calledWith(['from two.three import TestClass']));
+ done();
+ });
+ });
+
+ test('selection is added at line one', done => {
+ lookupSymbol.resolves([
+ new vscode.SymbolInformation(
+ 'TestClass', vscode.SymbolKind.Class, '', fileOneLoc
+ ),
+ ]);
+
+ showChoices.resolves('from one import TestClass');
+
+ autoimp.autoImport('TestClass').then(function() {
+ let line0 = vscode.window.activeTextEditor.document.lineAt(0);
+ assert.equal(line0.text, 'from one import TestClass');
+ done();
+ });
+ });
+});