Add PEP 604 style type hint support (NilsJPWerner#218)

src/parse/guess_type.ts

@@ -13,7 +13,7 @@ export function guessType(parameter: string): string {
 }
 
 function getTypeFromTyping(parameter: string): string {
-    const pattern = /\w+\s*:\s*(['"]?\w[\w\[\], \.]*['"]?)/;
+    const pattern = /\w+\s*:\s*(['"]?\w[\w\[\], |\.]*['"]?)/;
     const typeHint = pattern.exec(parameter);
 
     if (typeHint == null || typeHint.length !== 2) {

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/parse/tokenize_definition.ts

@@ -1,5 +1,5 @@
 export function tokenizeDefinition(functionDefinition: string): string[] {
-    const definitionPattern = /(?:def|class)\s+\w+\s*\(([\s\S]*)\)\s*(->\s*(["']?)[\w\[\], \.]*\3)?:\s*(?:#.*)?$/;
+    const definitionPattern = /(?:def|class)\s+\w+\s*\(([\s\S]*)\)\s*(->\s*(["']?)[\w\[\], |\.]*\3)?:\s*(?:#.*)?$/;
 
     const match = definitionPattern.exec(functionDefinition);
     if (match == undefined || match[1] == undefined) {
@@ -24,7 +24,7 @@ function tokenizeParameterString(parameterString: string): string[] {
 
     while (position >= 0) {
         const top = stack[stack.length - 1];
-        const char = parameterString.charAt(position);
+        let char = parameterString.charAt(position);
 
         /* todo
             '<' char,
@@ -81,6 +81,12 @@ function tokenizeParameterString(parameterString: string): string[] {
             case char === "\t" && stack.length === 0:
                 position -= 1;
                 continue;
+
+            // 9. Surround pipe character with whitespace
+            case char === "|":
+                char = ` ${char}`;
+                arg = ` ${arg}`;
+                break;
         }
 
         arg = char + arg;

src/test/integration/integration.test.ts

@@ -167,6 +167,17 @@ describe("Basic Integration Tests", function () {
             });
         });
 
+        it("generates a docstring using PEP 604 style type hints in file 7", async function () {
+            await testDocstringGeneration({
+                expectedOutputFilePath: path.resolve(
+                    __dirname,
+                    "./python_test_files/file_7_output.py",
+                ),
+                inputFilePath: path.resolve(__dirname, "./python_test_files/file_7.py"),
+                position: new vsc.Position(8, 0),
+            });
+        });
+
         it("generates a docstring for the starlark function", async function () {
             await testDocstringGeneration({
                 expectedOutputFilePath: path.resolve(

src/test/integration/python_test_files/file_7.py

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

src/test/integration/python_test_files/file_7_output.py

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

src/test/parse/guess_type.spec.ts

@@ -47,6 +47,13 @@ describe("guessType()", () => {
 
             expect(result2).to.equal("str");
         });
+
+        it("should get type from PEP 604 style type hint", () => {
+            const parameter = "arg: int | str";
+            const result = guessType(parameter);
+
+            expect(result).to.equal("int | str");
+        });
     });
 
     context("when the parameter is a kwarg", () => {

src/test/parse/parse_parameters.spec.ts

@@ -71,6 +71,15 @@ describe("parseParameters()", () => {
             });
         });
 
+        it("should parse PEP 604 type hint return types", () => {
+            const parameterTokens = ["-> int | float"];
+            const result = parseParameters(parameterTokens, [], "name");
+
+            expect(result.returns).to.deep.equal({
+                type: "int | float",
+            });
+        });
+
         it("should parse return types wrapped in single quotes", () => {
             expect(
                 parseParameters(["-> 'List[Type]'"], [], "name").returns

src/test/parse/tokenize_definition.spec.ts

@@ -67,6 +67,13 @@ describe("tokenizeDefinition()", () => {
         expect(result).to.have.ordered.members(["arg:string", "arg2:Callable[[], str]", "-> str"]);
     });
 
+    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 result = tokenizeDefinition(functionDefinition);
+
+        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", () => {
         const functionDefinition = "def func() -> str:";
         const result = tokenizeDefinition(functionDefinition);