Add return type hint completion

Author: njqdev

README.md

@@ -15,6 +15,6 @@ The extension is not published yet.
 
 ## Known Issues
 
-* The difference between classes and function calls when detecting types is determined by the first letter being upper case (unless the class or function is declared in the currently edited document).
+* The difference between classes and function calls when detecting types is determined by the first letter being upper case (unless the class or function is defined in the currently edited document).
 
 -------------------------------------------------------------------------------------------

src/completionProvider.ts

@@ -10,20 +10,20 @@ import {
     TextDocument
 } from "vscode";
 import { TypeResolver } from "./typeResolver";
-import { typeHintCharacter, Type } from "./syntax";
+import { paramHintTrigger, returnHintTrigger, Type } from "./syntax";
 
 /**
- * Provides one or more type hint {@link CompletionItem}.
+ * Provides one or more parameter type hint {@link CompletionItem}.
  */
-export class HintCompletionProvider implements CompletionItemProvider {
+export class ParamHintCompletionProvider implements CompletionItemProvider {
 
     public provideCompletionItems(
         doc: TextDocument, 
         pos: Position,
         token: CancellationToken,
         context: CompletionContext
     ): Thenable<CompletionList> | null {
-        if (context.triggerCharacter !== typeHintCharacter) {
+        if (context.triggerCharacter !== paramHintTrigger) {
             return null;
         }
         const items: CompletionItem[] = [];
@@ -59,4 +59,44 @@ export class HintCompletionProvider implements CompletionItemProvider {
         }
         return param;
     }
+}
+
+/**
+ * Provides one or more return type hint {@link CompletionItem}.
+ */
+export class ReturnHintCompletionProvider implements CompletionItemProvider {
+
+    public provideCompletionItems(
+        doc: TextDocument, 
+        pos: Position,
+        token: CancellationToken,
+        context: CompletionContext
+    ): Thenable<CompletionList> | null {
+        if (context.triggerCharacter !== returnHintTrigger) {
+            return null;
+        }
+        const items: CompletionItem[] = [];
+        const line = doc.lineAt(pos);
+
+        if (this.shouldProvideReturnHint(line, pos)) {         
+            pushDefaultCompletionItems(items);
+        }
+        return Promise.resolve(new CompletionList(items, false));
+    }
+
+    private shouldProvideReturnHint(line: TextLine, pos: Position): boolean {
+        if (pos.character > 0 && line.text.substr(pos.character - 2, 2) === "->") {
+            
+            if (new RegExp("^[*\t]*def.*\\) *->[: ]*$", "m").test(line.text)) {
+                return true;
+            }
+        }
+        return false;
+    }
+}
+
+function pushDefaultCompletionItems(items: CompletionItem[]) {
+    for (const type of Object.values(Type)) {
+        items.push(new CompletionItem(" " + type, CompletionItemKind.TypeParameter));
+    }
 }

src/extension.ts

@@ -1,6 +1,6 @@
 import * as vscode from 'vscode';
-import { HintCompletionProvider } from './completionProvider';
-import { typeHintCharacter } from "./syntax";
+import { ParamHintCompletionProvider, ReturnHintCompletionProvider } from './completionProvider';
+import { paramHintTrigger, returnHintTrigger } from "./syntax";
 
 // Called when the extension is activated.
 export function activate(context: vscode.ExtensionContext) {
@@ -11,7 +11,8 @@ export function activate(context: vscode.ExtensionContext) {
     }
 
     context.subscriptions.push(
-        vscode.languages.registerCompletionItemProvider('python', new HintCompletionProvider(), typeHintCharacter)
+        vscode.languages.registerCompletionItemProvider('python', new ParamHintCompletionProvider(), paramHintTrigger),
+        vscode.languages.registerCompletionItemProvider('python', new ReturnHintCompletionProvider(), returnHintTrigger)
     );
 }
 

src/syntax.ts

@@ -1,7 +1,8 @@
 
 export const anyTypeName: string = "[a-zA-Z_][a-zA-Z0-9_.]*";
 
-export const typeHintCharacter: string = ":";
+export const paramHintTrigger: string = ":";
+export const returnHintTrigger: string = ">";
 
 /**
  * Contains built-in Python types which can be hinted.