From b33bb1809e1f4721cba3f88b92bdc4aa32276ea7 Mon Sep 17 00:00:00 2001
From: Philip Nelson <philipnelson5@gmail.com>
Date: Tue, 2 Apr 2024 02:10:46 -0600
Subject: [PATCH] improve definition detection

reserved words 'class' and 'def' would break the identification of function and class definitions
---
 CHANGELOG.md                          |  1 +
 src/parse/get_definition.ts           | 11 +++++++----
 src/test/parse/get_definition.spec.ts | 20 ++++++++++++++++++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7dc72b8..f018e48 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## Unreleased
 
 -   Added doxygen format support ([#120](https://github.com/NilsJPWerner/autoDocstring/issues/120)) (@daluar @PhilipNelson5)
+-   Improve definition detection to allow reserved words in the default arguments ([#190](https://github.com/NilsJPWerner/autoDocstring/issues/190)) (@ylfzq @PhilipNelson5)
 
 [All Changes](https://github.com/NilsJPWerner/autoDocstring/compare/v0.6.1...master)
 
diff --git a/src/parse/get_definition.ts b/src/parse/get_definition.ts
index 5e6dd3c..919dc81 100644
--- a/src/parse/get_definition.ts
+++ b/src/parse/get_definition.ts
@@ -10,12 +10,15 @@ export function getDefinition(document: string, linePosition: number): string {
         return "";
     }
 
-    const pattern = /\b(((async\s+)?\s*def)|\s*class)\b/g;
+    // Match (async)? def name (...) and class name (...)
+    const pattern = /\b(((async\s+)?\s*def)|\s*class)\s+\w+\s*\(.*?\)/g;
 
-    // Get starting index of last def match in the preceding text
+    // Get starting index of last def or class match in the preceding text
     let index: number;
-    while (pattern.test(precedingText)) {
-        index = pattern.lastIndex - RegExp.lastMatch.length;
+    let match = pattern.exec(precedingText);
+    while (match) {
+        index = match.index;
+        match = pattern.exec(precedingText);
     }
 
     if (index == undefined) {
diff --git a/src/test/parse/get_definition.spec.ts b/src/test/parse/get_definition.spec.ts
index 7c4f3cb..972780b 100644
--- a/src/test/parse/get_definition.spec.ts
+++ b/src/test/parse/get_definition.spec.ts
@@ -14,6 +14,12 @@ describe("getDefinition()", () => {
             expect(result).to.equal("def basic_function(param1, param2 = abc):");
         });
 
+        it("should get a function definition that uses reserved words as default arguments", () => {
+            const result = getDefinition(functionWithReservedWordDefaults, 10);
+
+            expect(result).to.equal('def function_with_reserved_word_defaults( a: int, b="def", c="class" ) -> None:');
+        });
+
         it("should get an indented function definition", () => {
             const result = getDefinition(indentedFunction, 4);
 
@@ -84,6 +90,20 @@ def basic_function(param1, param2 = abc):
 def something_else():
 `;
 
+const functionWithReservedWordDefaults = `
+def another_func(a: str = "class, b: int): -> str
+    return a
+
+class Foo(object):
+    pass
+
+def function_with_reserved_word_defaults(
+    a: int, b="def", c="class"
+) -> None:
+    """
+    pass
+`;
+
 const indentedFunction = `
 Something Else