Add type guessing

Author: njqdev

.vscode/settings.json

@@ -7,5 +7,10 @@
         "out": true // set this to false to include "out" folder in search results
     },
     // Turn off tsc task auto detection since we have the necessary tasks as npm scripts
-    "typescript.tsc.autoDetect": "off"
+    "typescript.tsc.autoDetect": "off",
+    "[typescript]": {
+        "editor.rulers": [
+            120
+        ]
+    }
 }

package.json

@@ -1,7 +1,8 @@
 {
 	"name": "vscode-python-typehint",
 	"displayName": "Python Type Hint",
-	"description": "VSCode extension for Python type hint autocompletion.",
+    "description": "VSCode extension for Python type hint autocompletion.",
+    "publisher": "nj2532",
 	"repository": {
 		"type": "git",
 		"url": "https://github.com/nj2532/vscode-python-typehint.git"

src/completionProvider.ts

@@ -36,9 +36,7 @@ export class ParamHintCompletionProvider implements CompletionItemProvider {
             if (hint) {
                 items.push(new CompletionItem(" " + hint, CompletionItemKind.TypeParameter));
             } else {
-                for (const type of Object.values(Type)) {
-                    items.push(new CompletionItem(" " + type, CompletionItemKind.TypeParameter));
-                }
+                pushDefaultCompletionItems(items);
             }
         }
         return Promise.resolve(new CompletionList(items, false));
@@ -85,11 +83,10 @@ export class ReturnHintCompletionProvider implements CompletionItemProvider {
     }
 
     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 new RegExp("^[*\t]*def.*\\) *->[: ]*$", "m").test(line.text);
         }
         return false;
     }

src/typeResolver.ts

@@ -1,34 +1,78 @@
 import { TextDocument } from "vscode";
 import * as search from "./codeSearch";
+import { Type, returnHintTrigger } from "./syntax";
+import { ToTitleCase } from "./utils";
 
 export class TypeResolver {
 
+    private likelyTypes: Type[] = [ Type.String, Type.List, Type.Dict, Type.Bool, Type.Int, Type.Tuple, Type.Float ];
+
     /**
-     * Estimates the type of a parameter.
+     * Estimates the type of a parameter,
+     * by searching the document for an initialized variable with the same name.
+     * 
      * @param doc The document to search.
      * @param param The parameter name.
      */
     public EstimateType(doc: TextDocument, param: string): string | null {
         const documentText = doc.getText();
-        let type: string | null = "";
-        const match = new RegExp(`^[ \t]*${param} *=.`, "m").exec(documentText);
+        
+        let type = this.GetTypeIfEndsWith(param, "_", ...this.likelyTypes);
+        if (type) {
+            return type;
+        }
+        
+        const variableMatch = new RegExp(`^[ \t]*${param} *=.`, "m").exec(documentText);
 
-        if (match && match[0]) {
-            const valueStartPosition = doc.positionAt(match.index + match[0].length);
+        if (variableMatch && variableMatch[0]) {
+            const valueStartPosition = doc.positionAt(variableMatch.index + variableMatch[0].length);
             const line = doc.lineAt(valueStartPosition);
 
-            type = search.detectBasicType(line.text);
+            let typeName = search.detectBasicType(line.text);
 
-            if (type === null) {
-                type = search.detectNonBasicType(line.text, documentText);
+            if (typeName === null) {
+                typeName = search.detectNonBasicType(line.text, documentText);
             }
 
-            if (type !== null) {
-                if (!search.invalidTernaryOperator(type, line.text)) {
-                    return type;
+            if (typeName !== null) {
+                if (!search.invalidTernaryOperator(typeName, line.text)) {
+                    return typeName;
                 }
             }
         }
+        type = this.GuessType(param);
+        return type ? type : this.GetTypeIfEndsWith(param, "", Type.List, Type.Dict);
+    }
+
+    private GetTypeIfEndsWith(param: string, prefix: string, ...typesToCheckFor: Type[]): Type | null {
+
+        for (const type of typesToCheckFor) {
+            if (param.endsWith(`${prefix}${type}`)) {
+                return type;
+            }
+        }
         return null;
     }
+
+    private GuessType(param: string): Type | null {
+        if (param in this.TypeGuesses) {
+            return this.TypeGuesses[param];
+        }
+        return null;
+    }
+
+    /**
+     * Guesses used if a param with the same name is not found in the active document.
+     */
+    private TypeGuesses: { [key: string]: Type } = {
+        "string": Type.String,
+        "text": Type.String,
+        "path": Type.String,
+        "url": Type.String,
+        "uri": Type.String,
+        "fullpath": Type.String,
+        "full_path": Type.String,
+        "number": Type.Int,
+        "num": Type.Int
+    };
 }

src/utils.ts

@@ -0,0 +1,8 @@
+
+/**
+ * 
+ * @param s A string longer than 2 chars
+ */
+export function ToTitleCase(s: string) {
+    return `${s[0].toUpperCase}${s.substr(1)}`;
+}