Provide all typing hints if typing is not imported

Author: njqdev

src/typingHintProvider.ts

@@ -8,6 +8,7 @@ import { VariableSearchResult, TypeSearch } from "./typeSearch";
 export class TypingHintProvider {
 
     private typeContainer: DataTypeContainer;
+    private typingImportDetected: boolean = false;
     private fromTypingImport: boolean = false;
     private providedTypes: DataType[] = [];
     private typingPrefix: string = "typing";
@@ -28,27 +29,20 @@ export class TypingHintProvider {
      * @returns True if typing is imported.
      */
     public async detectTypingImport(docText: string): Promise<boolean> {
-        let m = new RegExp(
-            `^[ \t]*from typing import +([A-Z][a-zA-Z0-9 ,]+)`,
-            "m"
-        ).exec(docText);
-        
-        if (m) {
+
+        if (/^[ \t]*from typing import +([A-Z][a-zA-Z0-9 ,]+)/m.exec(docText)) {
             this.fromTypingImport = true;
-            return true;
+            this.typingImportDetected = true;
         } else {
-            m = new RegExp(
-                `^[ \t]*(import +typing +as +([a-zA-Z_][a-zA-Z0-9_-]*)|import +typing)`,
-                "m"
-            ).exec(docText);
-            if (m) {
-                if (m[2]) {
-                    this.typingPrefix = m[2];
+            const match = /^[ \t]*(import +typing +as +([a-zA-Z_][a-zA-Z0-9_-]*)|import +typing)/m.exec(docText);
+            if (match) {
+                if (match[2]) {
+                    this.typingPrefix = match[2];
                 }
-                return true;
+                this.typingImportDetected = true;
             }
         }
-        return false;
+        return this.typingImportDetected;
     }
 
     /**
@@ -61,7 +55,7 @@ export class TypingHintProvider {
         const type = this.typeContainer[typeName];
         if (type.category === TypeCategory.Collection) {
             this.providedTypes.push(type);
-            return this.toTypingString(type.name);
+            return this.toTypingString(type.name, this.fromTypingImport);
         }
         return null;
     }
@@ -78,7 +72,7 @@ export class TypingHintProvider {
             const type = this.typeContainer[searchResult.typeName];
             this.providedTypes.push(type);
 
-            const result: string[] = [ this.toTypingString(type.name)];
+            const result: string[] = [ this.toTypingString(type.name, this.fromTypingImport) ];
             let label = result[0];
 
             if (type.category === TypeCategory.Collection) {
@@ -97,7 +91,7 @@ export class TypingHintProvider {
                     if (this.typeContainer[elementType].name === PythonType.Dict) {
                         dictElementFound = true;
                     }
-                    label += this.toTypingString(elementType);
+                    label += this.toTypingString(elementType, this.fromTypingImport);
                     elementValue = elementValue.trim().substr(1);
                     elementType = TypeSearch.detectType(elementValue);
                     collectionCount++;
@@ -129,18 +123,33 @@ export class TypingHintProvider {
      * @returns An array of types.
      */
     public getRemainingHints(): string[] {
+        if (!this.typingImportDetected) {
+            return this.hintsForAllCollectionTypes();
+        }
         const result: string[] = [];
-
         for (const type of Object.values(this.typeContainer).filter(t => t.category === TypeCategory.Collection)) {
             if (!this.providedTypes.includes(type)) {
-                result.push(this.toTypingString(type.name));
+                result.push(this.toTypingString(type.name, this.fromTypingImport));
             }
         }
         return result;
     }
 
-    private toTypingString(typeName: string): string {
+    private hintsForAllCollectionTypes(): string[] {
+        const firstHalf: string[] = [];
+        const secondHalf: string[] = [];
+        for (const type of Object.values(this.typeContainer).filter(t => t.category === TypeCategory.Collection)) {
+            if (!this.providedTypes.includes(type)) {
+                const withoutPrefix = this.fromTypingImport || !this.typingImportDetected;
+                firstHalf.push(this.toTypingString(type.name, withoutPrefix));
+                secondHalf.push(this.toTypingString(type.name, !withoutPrefix));
+            }
+        }
+        return firstHalf.concat(secondHalf);
+    }
+
+    private toTypingString(typeName: string, withoutPrefix: boolean): string {
         const typingName = capitalized(typeName);
-        return this.fromTypingImport ? `${typingName}[` : `${this.typingPrefix}.${typingName}[`;
+        return withoutPrefix ? `${typingName}[` : `${this.typingPrefix}.${typingName}[`;
     }
 }

test/providers/typingHintProvider.test.ts

@@ -9,48 +9,54 @@ suite('TypingHintProvider', () => {
 
     const importTyping: string = "import x\nimport typing";
     const importTypingAsX: string = "import x\nimport typing as x";
-    const fromTypingImport: string = "import x\nfrom typing import List";
+    const fromTypingImport: string = "import x\nfrom typing import List, Dict, Tuple, Set";
     const typeContainer = getDataTypeContainer();
 
     suite('containsTyping', () => {
 
+        let provider: TypingHintProvider;
+
+        setup(() => {
+            provider = new TypingHintProvider(typeContainer);    
+        });
+
         test("returns true for 'import typing'", async () => {
-            const provider = new TypingHintProvider(typeContainer);
             const actual = await provider.detectTypingImport(importTyping);
             assert.equal(actual, true);
         });
 
         test("returns true for 'import typing as x'", async () => {
-            const provider = new TypingHintProvider(typeContainer);
             const actual = await provider.detectTypingImport(importTypingAsX);
             assert.equal(actual, true);
         });
 
         test("returns true for 'from typing import x'", async () => {
-            const provider = new TypingHintProvider(typeContainer);
             const actual = await provider.detectTypingImport(fromTypingImport);
             assert.equal(actual, true);
         });
     });
 
     suite("getHint", () => {
 
+        let provider: TypingHintProvider;
+
+        setup(() => {
+            provider = new TypingHintProvider(typeContainer);    
+        });
+
         test("returns typing.Type[", async () => {
-            const provider = new TypingHintProvider(typeContainer);
             const expected = "typing.List[";
-            providerTest(await provider.detectTypingImport(importTyping), getHintTest, provider, expected);
+            runTest(await provider.detectTypingImport(importTyping), getHintTest, provider, expected);
         });
 
         test("returns x.Type[ for 'import typing as x'", async () => {
-            const provider = new TypingHintProvider(typeContainer);
             const expected = "x.List[";
-            providerTest(await provider.detectTypingImport(importTypingAsX), getHintTest, provider, expected);
+            runTest(await provider.detectTypingImport(importTypingAsX), getHintTest, provider, expected);
         });
 
         test("returns Type[ for 'from typing' import", async () => {
-            const provider = new TypingHintProvider(typeContainer);
             const expected = "List[";
-            providerTest(await provider.detectTypingImport(fromTypingImport), getHintTest, provider, expected);
+            runTest(await provider.detectTypingImport(fromTypingImport), getHintTest, provider, expected);
         });
 
         function getHintTest(provider: TypingHintProvider, expected: string) {
@@ -71,55 +77,102 @@ suite('TypingHintProvider', () => {
         test("returns Type[ for empty collection", () => {
             const data = "[]";
             const expected = ["List["];
-            providerTest(typingImported, getTypingHints, provider, {data, expected }, PythonType.List);
+            runTest(typingImported, getHintsTest, provider, {data, expected }, PythonType.List);
         });
 
         test("returns Dict[ and 'Dict[key,' for dicts", () => {
             let data = "{ 1: 2 }";
             let expected = ["Dict[", "Dict[int"];
-            providerTest(typingImported, getTypingHints, provider, { data, expected }, PythonType.Dict);
+            runTest(typingImported, getHintsTest, provider, { data, expected }, PythonType.Dict);
         });
 
         test("handles nestled dicts", () => {
             let data = "[ { 1: 2 } ]";
             let expected = ["List[", "List[Dict[int"];
-            providerTest(typingImported, getTypingHints, provider, { data, expected }, PythonType.List);
+            runTest(typingImported, getHintsTest, provider, { data, expected }, PythonType.List);
         });
 
         test("returns Type[ and Type[type] for non-dicts", () => {
             let data = "['str']";
             let expected = ["List[", "List[str]"];
-            providerTest(typingImported, getTypingHints, provider, { data, expected }, PythonType.List);
+            runTest(typingImported, getHintsTest, provider, { data, expected }, PythonType.List);
 
             data = "(1, {'ignore': 'this'})";
             expected = ["Tuple[", "Tuple[int]"];
-            providerTest(typingImported, getTypingHints, provider, { data, expected }, PythonType.Tuple);
+            runTest(typingImported, getHintsTest, provider, { data, expected }, PythonType.Tuple);
 
             data = "{ 1, 2 }";
             expected = ["Set[", "Set[int]"];
-            providerTest(typingImported, getTypingHints, provider, { data, expected }, PythonType.Set);
+            runTest(typingImported, getHintsTest, provider, { data, expected }, PythonType.Set);
         });
 
         test("adds typing prefixes for 'import typing' imports", async () => {
             let p = new TypingHintProvider(typeContainer);
             let data = "[ { 1: 2 } ]";
             let expected = ["typing.List[", "typing.List[typing.Dict[int"];
-            providerTest(
+            runTest(
                 await p.detectTypingImport(importTyping),
-                getTypingHints,
+                getHintsTest,
                 p,
                 { data, expected },
                 PythonType.List
             );
         });
 
-        function getTypingHints(provider: TypingHintProvider, testCase: TestCase, type: PythonType) {
+        function getHintsTest(provider: TypingHintProvider, testCase: TestCase, type: PythonType) {
             const actual = provider.getHints(varSearchResult(type, testCase.data));
             assert.deepEqual(actual, testCase.expected, messageFor(testCase.data, testCase.expected, actual));
         }
     });
 
-    function providerTest(typingDetected: boolean, test: (...params: any) => void, ...args: any) {
+    suite("getRemainingHints", () => {
+
+        let provider: TypingHintProvider;
+
+        setup(() => {
+            provider = new TypingHintProvider(typeContainer);    
+        });
+
+        test("returns all typing hints if typing is not imported, without prefix followed by with prefix", async () => {
+            await provider.detectTypingImport("");
+            const expected = typingTypes(false).concat(typingTypes(true));
+            const actual = provider.getRemainingHints();
+            assert.deepEqual(actual, expected);
+        });
+
+        test("returns hints without prefix first, for from typing import", async () => {
+            const expected = typingTypes(false);
+            runTest(
+                await provider.detectTypingImport(fromTypingImport),
+                getRemainingHintsTest,
+                provider,
+                expected
+            );
+        });
+
+        test("returns hints with prefix first, if imported", async () => {
+            const expected = typingTypes(true);
+            runTest(
+                await provider.detectTypingImport(importTyping),
+                getRemainingHintsTest,
+                provider,
+                expected
+            );
+        });
+
+        
+        function getRemainingHintsTest(provider: TypingHintProvider, expected: string) {
+            const actual = provider.getRemainingHints();
+            assert.deepEqual(actual, expected);
+        }
+    });
+
+    const typingTypes = (withPrefix: boolean) => {
+        const prefix = withPrefix ? "typing." : "";
+        return [`${prefix}Dict[`, `${prefix}List[`, `${prefix}Set[`, `${prefix}Tuple[` ];
+    };
+
+    function runTest(typingDetected: boolean, test: (...params: any) => void, ...args: any) {
         if (typingDetected) {
             test(...args);
         } else {