CodeLens for Shebangs

package.json

@@ -51,6 +51,7 @@
     "onCommand:python.runtests",
     "onCommand:python.debugtests",
     "onCommand:python.setInterpreter",
+    "onCommand:python.setShebangInterpreter",
     "onCommand:python.viewTestUI",
     "onCommand:python.viewTestOutput",
     "onCommand:python.selectAndRunTestMethod",

src/client/extension.ts

@@ -7,6 +7,7 @@ import { PythonDefinitionProvider } from './providers/definitionProvider';
 import { PythonReferenceProvider } from './providers/referenceProvider';
 import { PythonRenameProvider } from './providers/renameProvider';
 import { PythonFormattingEditProvider } from './providers/formatProvider';
+import { ShebangCodeLensProvider } from './providers/shebangCodeLensProvider'
 import * as sortImports from './sortImports';
 import { LintProvider } from './providers/lintProvider';
 import { PythonSymbolProvider } from './providers/symbolProvider';
@@ -105,6 +106,7 @@ export async function activate(context: vscode.ExtensionContext) {
     context.subscriptions.push(vscode.languages.registerHoverProvider(PYTHON, new PythonHoverProvider(context, jediProx)));
     context.subscriptions.push(vscode.languages.registerReferenceProvider(PYTHON, new PythonReferenceProvider(context, jediProx)));
     context.subscriptions.push(vscode.languages.registerCompletionItemProvider(PYTHON, new PythonCompletionItemProvider(context, jediProx), '.'));
+    context.subscriptions.push(vscode.languages.registerCodeLensProvider(PYTHON, new ShebangCodeLensProvider()))
 
     const symbolProvider = new PythonSymbolProvider(context, jediProx);
     context.subscriptions.push(vscode.languages.registerDocumentSymbolProvider(PYTHON, symbolProvider));

src/client/providers/setInterpreterProvider.ts

@@ -4,6 +4,7 @@ import * as vscode from 'vscode';
 import * as settings from './../common/configSettings';
 import { InterpreterManager } from '../interpreter';
 import { PythonInterpreter } from '../interpreter/contracts';
+import { ShebangCodeLensProvider } from './shebangCodeLensProvider';
 
 
 interface PythonPathQuickPickItem extends vscode.QuickPickItem {
@@ -14,6 +15,7 @@ export class SetInterpreterProvider implements vscode.Disposable {
     private disposables: vscode.Disposable[] = [];
     constructor(private interpreterManager: InterpreterManager) {
         this.disposables.push(vscode.commands.registerCommand("python.setInterpreter", this.setInterpreter.bind(this)));
+        this.disposables.push(vscode.commands.registerCommand("python.setShebangInterpreter", this.setShebangInterpreter.bind(this)));
     }
     public dispose() {
         this.disposables.forEach(disposable => disposable.dispose());
@@ -60,4 +62,11 @@ export class SetInterpreterProvider implements vscode.Disposable {
     private setInterpreter() {
         this.presentQuickPick();
     }
-}
+
+    private setShebangInterpreter() {
+        const shebang = ShebangCodeLensProvider.detectShebang(vscode.window.activeTextEditor.document);
+        if (shebang) {
+            this.interpreterManager.setPythonPath(shebang);
+        }
+    }
+}

src/client/providers/shebangCodeLensProvider.ts

@@ -0,0 +1,57 @@
+"use strict";
+import * as vscode from 'vscode'
+import { TextDocument, CodeLens, CancellationToken } from 'vscode'
+
+export class ShebangCodeLensProvider implements vscode.CodeLensProvider {
+    private settings;
+
+    // reload codeLenses on every configuration change.
+    onDidChangeCodeLenses: vscode.Event<void> = vscode.workspace.onDidChangeConfiguration;
+
+    public provideCodeLenses(document: TextDocument, token: CancellationToken): Thenable<CodeLens[]> {
+        this.settings = vscode.workspace.getConfiguration('python');
+        const codeLenses = this.createShebangCodeLens(document);
+
+        return Promise.resolve(codeLenses);
+    }
+
+    private createShebangCodeLens(document: TextDocument) {
+        const shebang = ShebangCodeLensProvider.detectShebang(document)
+        if (!shebang || shebang === this.settings.get('pythonPath')) { 
+            // no shebang detected or interpreter is already set to shebang
+            return; 
+        }
+
+        // create CodeLens
+        const firstLine = document.lineAt(0);
+        const startOfShebang = new vscode.Position(0, 0);
+        const endOfShebang = new vscode.Position(0, firstLine.text.length - 1);
+        const shebangRange = new vscode.Range(startOfShebang, endOfShebang);
+
+        const cmd : vscode.Command = { 
+            command: 'python.setShebangInterpreter', 
+            title: 'Set as interpreter'
+        }
+
+        const codeLenses = [(new CodeLens(shebangRange, cmd))];
+        return codeLenses;
+    }
+
+    public static detectShebang(document: TextDocument) {
+        let error = false;
+
+        let firstLine = document.lineAt(0);
+        if (firstLine.isEmptyOrWhitespace) {
+            error = true;
+        }
+
+        if (!error && "#!" === firstLine.text.substr(0, 2)) {
+            // Shebang detected
+            const shebang = firstLine.text.substr(2).trim();
+            return shebang;
+        }
+
+        return null;
+    }
+
+}

src/test/providers/shebangCodeLenseProvider.test.ts

@@ -0,0 +1,69 @@
+import * as assert from 'assert';
+import * as path from 'path';
+import * as vscode from 'vscode';
+import { ShebangCodeLensProvider } from '../../client/providers/shebangCodeLensProvider'
+
+import { initialize, IS_TRAVIS, closeActiveWindows } from '../initialize';
+
+const autoCompPath = path.join(__dirname, '..', '..', '..', 'src', 'test', 'pythonFiles', 'shebang');
+const fileShebang = path.join(autoCompPath, 'shebang.py');
+const filePlain = path.join(autoCompPath, 'plain.py');
+
+var settings = vscode.workspace.getConfiguration("python");
+const origPythonPath = settings.get("pythonPath");
+
+suite("Shebang detection", () => {
+    suiteSetup(async () => {
+        await initialize();
+    });
+
+    suiteTeardown(async () => {
+        await vscode.workspace.getConfiguration("python").update("pythonPath", origPythonPath);
+    });
+
+    teardown(() => closeActiveWindows());
+    setup(() => {
+        settings = vscode.workspace.getConfiguration("python");
+    });
+
+    test("Shebang available, CodeLens showing", async () => {
+        await settings.update("pythonPath", "python");
+        const editor = await openFile(fileShebang);
+        const codeLenses = await setupCodeLens(editor);
+
+        assert.equal(codeLenses.length, 1, "No CodeLens available");
+        let codeLens = codeLenses[0];
+        assert(codeLens.range.isSingleLine, 'Invalid CodeLens Range');
+        assert.equal(codeLens.command.command, 'python.setShebangInterpreter');
+
+    });
+
+    test("Shebang available, CodeLens hiding", async () => {
+        await settings.update("pythonPath", "/usr/bin/test");
+        const editor = await openFile(fileShebang);
+        const codeLenses = await setupCodeLens(editor);
+        assert(!codeLenses, "CodeLens available although interpreters are equal");
+
+    });
+
+    test("Shebang missing, CodeLens hiding", async () => {
+        const editor = await openFile(filePlain);
+        const codeLenses = await setupCodeLens(editor);
+        assert(!codeLenses, "CodeLens available although no shebang");
+
+    });
+
+    async function openFile(fileName: string) {
+        const document = await vscode.workspace.openTextDocument(fileName);
+        const editor = await vscode.window.showTextDocument(document);
+        assert(vscode.window.activeTextEditor, 'No active editor');
+        return editor;
+    }
+
+    async function setupCodeLens(editor: vscode.TextEditor) {
+        const document = editor.document;
+        const codeLensProvider = new ShebangCodeLensProvider();
+        const codeLenses = await codeLensProvider.provideCodeLenses(document, null);
+        return codeLenses;
+    }
+});

src/test/pythonFiles/shebang/plain.py

@@ -0,0 +1,2 @@
+
+print("dummy")

src/test/pythonFiles/shebang/shebang.py

@@ -0,0 +1,3 @@
+#!/usr/bin/test
+
+print("dummy")