Multiple estimations, shift responsibility to TypeResolver

Author: njqdev

src/codeSearch.ts

@@ -1,10 +1,10 @@
-import {
-    TextLine, 
-    Position,
-    TextDocument
- } from "vscode";
+import { anyTypeName, Types } from "./python";
 
-import { anyTypeName, Type } from "./syntax";
+
+export interface TypeSearchResult {
+    typeName: string | null;
+        
+}
 
 /**
  * Detects the type of an initialized variable.
@@ -15,8 +15,8 @@ import { anyTypeName, Type } from "./syntax";
  */
 export function detectBasicType(src: string, srcIsLineOfCode = true): string | null {
 
-    for (const typeName of Object.values(Type)) {
-        let r = getTypeRegEx(typeName, srcIsLineOfCode ? "= *" : "");
+    for (const typeName of typeSearchOrder) {
+        let r = typeSearchRegExp(typeName, srcIsLineOfCode ? "= *" : "");
         if (r.test(src)) {
             return typeName;
         }
@@ -41,7 +41,7 @@ export function detectNonBasicType(lineText: string, documentText: string): stri
 
     if (match[0].endsWith("(")) {
 
-        if (isClass(match[1], documentText)) {
+        if (classWithSameName(match[1], documentText)) {
             return match[1];
         }
 
@@ -87,10 +87,7 @@ export function detectNonBasicType(lineText: string, documentText: string): stri
  */
 export function invalidTernaryOperator(typeName: string, lineSrc: string) {
 
-    const regExp = new RegExp(
-        " if +[^ ]+ +else( +[^ ]+) *$", 
-        "m"
-    );
+    const regExp = new RegExp(" if +[^ ]+ +else( +[^ ]+) *$", "m");
 
     let ternaryMatch = regExp.exec(lineSrc);
     while (ternaryMatch) {
@@ -109,11 +106,14 @@ export function invalidTernaryOperator(typeName: string, lineSrc: string) {
     return false;
 }
 
-function isClass(object: string, documentText: string) {
-    return new RegExp(
-        `^ *class +${object}`,
-        "m"
-    ).test(documentText);
+/**
+ * Searches for a class with the same name as object and returns the name if found.
+ * @param object The object.
+ * @param documentText The text to search
+ */
+export function classWithSameName(object: string, documentText: string): string | null {
+    const clsMatch = new RegExp(`^ *class +(${object})`, "mi").exec(documentText);
+    return clsMatch ? clsMatch[1] : null;
 }
 
 function importFound(object: string, documentText: string): boolean {
@@ -128,37 +128,51 @@ function isProbablyAClass(lineText: string): boolean {
 }
 
 function isType(text: string): boolean {
-    return Object.values(Type).includes(text as Type);
+    return Object.values(Types).includes(text as Types);
 }
 
+const typeSearchOrder = [
+    Types.List, 
+    Types.Bool,
+    Types.Complex,
+    Types.Float,
+    Types.String,   
+    Types.Tuple,
+    Types.Set,
+    Types.Dict,
+    Types.Int,
+    Types.Object
+];
+
 /**
  * Get a new RegExp for finding basic types and {@class object}.
  * 
  * @param typeName the type name
  * @param prefix a prefix added to the RegExp pattern
  */
-function getTypeRegEx(typeName: string, prefix: string): RegExp {
+function typeSearchRegExp(typeName: string, prefix: string): RegExp {
     switch (typeName) {
-        case Type.Bool:
-            return new RegExp(`${prefix}(True|False|bool\\()`, "m");
-        case Type.Dict:
-            return new RegExp(`${prefix}({|dict\\()`, "m");
-        case Type.Int:
-            return new RegExp(`${prefix}(-*[0-9]+(?!(\\.| *\\)| *,))|int\\()`, "m");
-        case Type.List:
+        case Types.List:
             return new RegExp(`${prefix}(\\[|list\\()`, "m");
-        case Type.String:
+        case Types.Bool:
+            return new RegExp(`${prefix}(True|False|bool\\()`, "m");
+        case Types.Complex:
+            return new RegExp(`${prefix}(\\(complex\\(|[[0-9+*\\/ -.]*[0-9][jJ])`, "m");
+        case Types.Float:
+            return new RegExp(`${prefix}(-*[0-9+*\/ -]*\\.[0-9]|float\\()`, "m");
+        case Types.Tuple:
+            return new RegExp(`${prefix}(\\(|tuple\\()`, "m"); 
+        case Types.String:
             return new RegExp(`${prefix}(['\"]{3}|(\\( *)?\"[^\"]*\"(?! *,)|(\\( *)?'[^']*'(?! *,)|str\\()`, "m");
-        case Type.Float:
-            return new RegExp(`${prefix}(-*[0-9]*\\.[0-9]+|float\\()`, "m");
-        case Type.Tuple:
-            return new RegExp(`${prefix}(\\(([^'\",)]+,|\"[^\"]*\"(?= *,)|'[^']*'(?= *,))|tuple\\()`, "m"); 
-        case Type.Complex:
-            return new RegExp(`${prefix}(\\(complex\\(|[0-9][0-9+-/*.]*[jJ])`, "m");
-        case Type.Object:
+        case Types.Set:
+            return new RegExp(`${prefix}({[^:]+[,}]|set\\()`, "m");
+        case Types.Dict:
+            return new RegExp(`${prefix}({|dict\\()`, "m");
+        case Types.Int:
+            return new RegExp(`${prefix}(-*[0-9]|int\\()`, "m");
+        case Types.Object:
             return new RegExp(`${prefix}object\\(`, "m");        
         default:
             return new RegExp(`^.*$`, "m");
     }
 }
-

src/completionProvider.ts

@@ -9,13 +9,27 @@ import {
     TextLine,
     TextDocument
 } from "vscode";
-import { TypeResolver } from "./typeResolver";
-import { paramHintTrigger, returnHintTrigger, Type } from "./syntax";
+import { TypeResolver, TypeResolution } from "./typeResolver";
+import { paramHintTrigger, returnHintTrigger, Types } from "./python";
+
+
+abstract class CompletionProvider {
+
+    protected pushTypesToItems(types: string[], completionItems: CompletionItem[]) {
+        for (const type of types) {
+            const item = new CompletionItem(" " + type, CompletionItemKind.TypeParameter);
+    
+            // Prioritize type estimations and sort remaining items alphabetically
+            item.sortText = `999${type}`;   
+            completionItems.push(item);
+        }
+    }
+}
 
 /**
  * Provides one or more parameter type hint {@link CompletionItem}.
  */
-export class ParamHintCompletionProvider implements CompletionItemProvider {
+export class ParamHintCompletionProvider extends CompletionProvider implements CompletionItemProvider {
 
     public provideCompletionItems(
         doc: TextDocument, 
@@ -31,13 +45,10 @@ export class ParamHintCompletionProvider implements CompletionItemProvider {
         const param = this.findParam(line, pos);
 
         if (param && param.length > 0) {         
-            let hint = new TypeResolver().EstimateType(doc, param);
-
-            if (hint) {
-                items.push(new CompletionItem(" " + hint, CompletionItemKind.TypeParameter));
-            } else {
-                pushDefaultCompletionItems(items);
-            }
+            const resolution = new TypeResolver().ResolveTypes(param, doc);
+            this.pushTypeResolutionToItems(resolution, items);
+        } else {
+            this.pushTypesToItems(Object.values(Types), items);
         }
         return Promise.resolve(new CompletionList(items, false));
     }
@@ -57,12 +68,27 @@ export class ParamHintCompletionProvider implements CompletionItemProvider {
         }
         return param;
     }
+    
+    private pushTypeResolutionToItems(resolution: TypeResolution, items: CompletionItem[]) {
+
+        if (resolution.estimations) {
+            for (let i = 0; i < resolution.estimations.length; i++) {
+                const typeName = resolution.estimations[i];
+                const item = new CompletionItem(" " + typeName, CompletionItemKind.TypeParameter);
+                item.sortText = `${i}${typeName}`;
+                item.preselect = true;
+                items.push(item);
+            }       
+
+        }
+        this.pushTypesToItems(resolution.remainingTypes, items);
+    }
 }
 
 /**
  * Provides one or more return type hint {@link CompletionItem}.
  */
-export class ReturnHintCompletionProvider implements CompletionItemProvider {
+export class ReturnHintCompletionProvider extends CompletionProvider implements CompletionItemProvider {
 
     public provideCompletionItems(
         doc: TextDocument, 
@@ -77,7 +103,7 @@ export class ReturnHintCompletionProvider implements CompletionItemProvider {
         const line = doc.lineAt(pos);
 
         if (this.shouldProvideReturnHint(line, pos)) {         
-            pushDefaultCompletionItems(items);
+            this.pushTypesToItems(Object.values(Types), items);
         }
         return Promise.resolve(new CompletionList(items, false));
     }
@@ -90,10 +116,4 @@ export class ReturnHintCompletionProvider implements CompletionItemProvider {
         }
         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 { ParamHintCompletionProvider, ReturnHintCompletionProvider } from './completionProvider';
-import { paramHintTrigger, returnHintTrigger } from "./syntax";
+import { paramHintTrigger, returnHintTrigger } from "./python";
 
 // Called when the extension is activated.
 export function activate(context: vscode.ExtensionContext) {

src/python.ts

@@ -0,0 +1,76 @@
+
+export const anyTypeName: string = "[a-zA-Z_][a-zA-Z0-9_.]*";
+
+export const paramHintTrigger: string = ":";
+export const returnHintTrigger: string = ">";
+
+
+export class DataType {
+    name: Types;
+    category: TypeCategory;
+
+    constructor(name: Types, category: TypeCategory) {
+        this.name = name;
+        this.category = category;
+    }
+}
+
+export interface DataTypes {
+     [key: string]: DataType 
+};
+
+export const getDataTypes = (): DataTypes => {
+    return {
+        bool: new DataType(Types.Bool, typeCategories.bool),
+        complex: new DataType(Types.Complex, typeCategories.complex),
+        dict: new DataType(Types.Dict, typeCategories.dict),
+        float: new DataType(Types.Float, typeCategories.float),
+        int: new DataType(Types.Int, typeCategories.int),
+        list: new DataType(Types.List, typeCategories.list),
+        object: new DataType(Types.Object, typeCategories.object),
+        set: new DataType(Types.Set, typeCategories.set),
+        str: new DataType(Types.String, typeCategories.string),
+        tuple: new DataType(Types.Tuple, typeCategories.tuple)
+    };
+};
+
+/**
+ * Names of built-in Python types which can be hinted. 
+ */
+export enum Types {
+    Bool = "bool",
+    Complex = "complex",
+    Dict = "dict",
+    Float = "float",
+    Int = "int",
+    List = "list",
+    Object = "object",
+    Set = "set",
+    String = "str",
+    Tuple = "tuple",
+}
+
+/**
+ * Categories of Python types.
+ */
+export enum TypeCategory {
+    Abstract,
+    Basic,
+    Collection
+}
+
+/**
+ * Type name keys with Type values. 
+ */
+const typeCategories: { [key: string]: TypeCategory } = {
+    bool: TypeCategory.Basic,
+    complex: TypeCategory.Basic,
+    dict: TypeCategory.Collection,
+    float: TypeCategory.Basic,
+    int: TypeCategory.Basic,
+    list: TypeCategory.Collection,
+    object: TypeCategory.Abstract,
+    set: TypeCategory.Collection,
+    string: TypeCategory.Basic,
+    tuple: TypeCategory.Collection
+};

src/syntax.ts

@@ -60,6 +60,10 @@ suite('detectBasicType', function() {
         src = "var = -.3";
         actual = detectBasicType(src);
         assert.equal(actual, expected, getErrorMessage(src, expected, actual));
+
+        src = "var = 1 + 2 - 1 * 2 /  2.0";
+        actual = detectBasicType(src);
+        assert.equal(actual, expected, getErrorMessage(src, expected, actual));
     });
 
     test("detects complex numbers", () => {
@@ -126,6 +130,22 @@ suite('detectBasicType', function() {
         let actual = detectBasicType(src);
         assert.equal(actual, expected);
     });
+
+    test("detects sets", () => {
+        const expected = "set";
+
+        let src = "var = {'dont return dict please'}";
+        let actual = detectBasicType(src);
+        assert.equal(actual, expected);
+
+        src = "var = {1, 2}";
+        actual = detectBasicType(src);
+        assert.equal(actual, expected, getErrorMessage(src, expected, actual));
+
+        src = "var = {1 , 2}";
+        actual = detectBasicType(src);
+        assert.equal(actual, expected, getErrorMessage(src, expected, actual));
+    });
 
     test("detects type() call", () => {
         const testCases: TestCase[] = [
@@ -134,7 +154,8 @@ suite('detectBasicType', function() {
             { data: "var = list(foo)", expected: "list"},
             { data: "var = dict(foo)", expected: "dict"},
             { data: "var = tuple(foo)", expected: "tuple"},
-            { data: "var = str(1)", expected: "str"}
+            { data: "var = str(1)", expected: "str"},
+            { data: "var = set([1])", expected: "set"}
         ];
         for (const c of testCases) {
             assert.equal(detectBasicType(c.data), c.expected);