Fixes #212: Parse Literal type hints in function definition (#223)

src/parse/guess_type.ts

@@ -2,7 +2,7 @@ export function guessType(parameter: string): string {
     parameter = parameter.trim();
 
     if (hasTypeHint(parameter)) {
-        return getTypeFromTyping(parameter);
+        return getTypeFromTypeHint(parameter);
     }
 
     if (isKwarg(parameter)) {
@@ -12,15 +12,25 @@ export function guessType(parameter: string): string {
     return guessTypeFromName(parameter);
 }
 
-function getTypeFromTyping(parameter: string): string {
-    const pattern = /\w+\s*:\s*(['"]?\w[\w\[\], |\.]*['"]?)/;
-    const typeHint = pattern.exec(parameter);
+function getTypeFromTypeHint(parameter: string): string {
+    const sections = parameter.split(":");
+    if (sections.length !== 2) {
+        return undefined;
+    }
+    const typeHint = sections[1];
 
-    if (typeHint == null || typeHint.length !== 2) {
+    const pattern = /\s*(['"]?\w["'\w\[\], |\.]*['"]?)/;
+    const typeHintRegex = pattern.exec(typeHint);
+
+    if (typeHintRegex == null) {
         return undefined;
     }
 
-    return typeHint[1].replace(/['"]+/g, '').trim();
+    // Remove enclosing quotes
+    let type = typeHintRegex[0].trim();
+    type = type.replace(/^['"]|['"]$/g, "");
+
+    return type;
 }
 
 function guessTypeFromDefaultValue(parameter: string): string {

src/parse/parse_parameters.ts

@@ -115,7 +115,7 @@ function parseYields(parameters: string[], body: string[]): Yields {
 }
 
 function parseReturnFromDefinition(parameters: string[]): Returns | null {
-    const pattern = /^->\s*(["']?)([\w\[\], |\.]*)\1/;
+    const pattern = /^->\s*(["']?)(['"\w\[\], |\.]*)\1/;
 
     for (const param of parameters) {
         const match = param.trim().match(pattern);

src/test/integration/integration.test.ts

@@ -141,7 +141,7 @@ describe("Basic Integration Tests", function () {
                     "./python_test_files/file_4_output.py",
                 ),
                 inputFilePath: path.resolve(__dirname, "./python_test_files/file_4.py"),
-                position: new vsc.Position(5, 0),
+                position: new vsc.Position(6, 0),
             });
         });
 

src/test/integration/python_test_files/file_4.py

@@ -1,3 +1,4 @@
+from typing import Callable, List, Enum
 
 def build(self,
             my_id: uuid.UUID,

src/test/integration/python_test_files/file_4_output.py

@@ -1,3 +1,4 @@
+from typing import Callable, List, Enum
 
 def build(self,
             my_id: uuid.UUID,

src/test/integration/python_test_files/file_7.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
-
+from typing import List, Dict
 
 def function(
     arg1: int,
-    arg2: list[str] | dict[str, int] | Thing,
+    arg2: List[str] | Dict[str, int] | None,
     kwarg1: int | float = 1
-) -> list[str] | dict[str, int] | Thing:
+) -> List[str] | Dict[str, int] | None:
 
     return arg2

src/test/integration/python_test_files/file_7_output.py

@@ -1,20 +1,20 @@
 from __future__ import annotations
-
+from typing import List, Dict
 
 def function(
     arg1: int,
-    arg2: list[str] | dict[str, int] | Thing,
+    arg2: List[str] | Dict[str, int] | None,
     kwarg1: int | float = 1
-) -> list[str] | dict[str, int] | Thing:
+) -> List[str] | Dict[str, int] | None:
     """_summary_
 
     :param arg1: _description_
     :type arg1: int
     :param arg2: _description_
-    :type arg2: list[str] | dict[str, int] | Thing
+    :type arg2: List[str] | Dict[str, int] | None
     :param kwarg1: _description_, defaults to 1
     :type kwarg1: int | float, optional
     :return: _description_
-    :rtype: list[str] | dict[str, int] | Thing
+    :rtype: List[str] | Dict[str, int] | None
     """
     return arg2

src/test/parse/guess_type.spec.ts

@@ -54,6 +54,20 @@ describe("guessType()", () => {
 
             expect(result).to.equal("int | str");
         });
+
+        it("should get type from args with Literal type hints", () => {
+            const parameter = `param1: Literal["foo"]`;
+            const result = guessType(parameter);
+
+            expect(result).to.equal(`Literal["foo"]`);
+        });
+
+        it("should get type from kwargs with Literal type hint", () => {
+            const parameter = `param2: Literal['bar'] = "bar"`;
+            const result = guessType(parameter);
+
+            expect(result).to.equal(`Literal['bar']`);
+        });
     });
 
     context("when the parameter is a kwarg", () => {

src/test/parse/parse_parameters.spec.ts

@@ -51,6 +51,25 @@ describe("parseParameters()", () => {
         ]);
     });
 
+    it("should parse args, kwargs, and return with Literal type hints", () => {
+        const parameterTokens = [
+            `param1: Literal["foo"]`,
+            `param2: Literal['bar'] = "bar"`,
+            `-> Literal["baz"]`,
+        ];
+        const result = parseParameters(parameterTokens, [], "name");
+
+        expect(result).to.eql({
+            name: "name",
+            args: [{ var: "param1", type: `Literal["foo"]` }],
+            kwargs: [{ var: "param2", default: `"bar"`, type: `Literal['bar']` }],
+            returns: { type: `Literal["baz"]` },
+            decorators: [],
+            exceptions: [],
+            yields: undefined,
+        });
+    });
+
     it("should parse kwargs with and without type hints", () => {
         const parameterTokens = ["param1: List[int] = [1,2]", "param2 = 'abc'"];
         const result = parseParameters(parameterTokens, [], "name");
@@ -81,18 +100,14 @@ describe("parseParameters()", () => {
         });
 
         it("should parse return types wrapped in single quotes", () => {
-            expect(
-                parseParameters(["-> 'List[Type]'"], [], "name").returns
-            ).to.deep.equal({
-                type: "List[Type]"
+            expect(parseParameters(["-> 'List[Type]'"], [], "name").returns).to.deep.equal({
+                type: "List[Type]",
             });
         });
 
         it("should parse return types wrapped in double quotes", () => {
-            expect(
-                parseParameters(['-> "List[Type]"'], [], "name").returns
-            ).to.deep.equal({
-                type: "List[Type]"
+            expect(parseParameters(['-> "List[Type]"'], [], "name").returns).to.deep.equal({
+                type: "List[Type]",
             });
         });
 

src/test/parse/tokenize_definition.spec.ts

@@ -68,10 +68,15 @@ describe("tokenizeDefinition()", () => {
     });
 
     it("should tokenize pep604 parameter and return types", () => {
-        const functionDefinition = "def func(arg: int | float | str, arg2: dict[str, str] | list[str]) -> int | float | str:";
+        const functionDefinition =
+            "def func(arg: int | float | str, arg2: dict[str, str] | list[str]) -> int | float | str:";
         const result = tokenizeDefinition(functionDefinition);
 
-        expect(result).to.have.ordered.members(["arg:int | float | str", "arg2:dict[str, str] | list[str]", "-> int | float | str"]);
+        expect(result).to.have.ordered.members([
+            "arg:int | float | str",
+            "arg2:dict[str, str] | list[str]",
+            "-> int | float | str",
+        ]);
     });
 
     it("should tokenize pep484 return types", () => {
@@ -81,6 +86,13 @@ describe("tokenizeDefinition()", () => {
         expect(result).to.have.ordered.members(["-> str"]);
     });
 
+    it("should tokenize Literal types", () => {
+        const functionDefinition = `def func(x: Literal["r", "w", "a"]) -> Literal[1, 2, 3]:`;
+        const result = tokenizeDefinition(functionDefinition);
+
+        expect(result).to.have.ordered.members([`x:Literal["r", "w", "a"]`, `-> Literal[1, 2, 3]`]);
+    });
+
     it("should split class definition arguments", () => {
         const functionDefinition = "class abc_c(arg, arg_2):";
         const result = tokenizeDefinition(functionDefinition);