From c79fda8cbdf80ea6dbadbf8803804ac1e5f959f5 Mon Sep 17 00:00:00 2001
From: Clif Bratcher <inno@users.noreply.github.com>
Date: Fri, 7 Jun 2024 09:28:03 -0400
Subject: [PATCH] Fixes #3 - Incorrect handling for first comment after
 function definition (#13)

* Use ast to detect the last line of the last argument, stop iterating once we're past it
---
 simplecli/simplecli.py     | 15 ++++++++++++++-
 tests/test_extract_args.py | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/simplecli/simplecli.py b/simplecli/simplecli.py
index 463695d..f64b159 100644
--- a/simplecli/simplecli.py
+++ b/simplecli/simplecli.py
@@ -1,10 +1,12 @@
 from __future__ import annotations
+import ast
 import contextlib
 import inspect
 import io
 import os
 import re
 import sys
+import textwrap
 from collections import OrderedDict
 from collections.abc import Generator
 from tokenize import (
@@ -482,15 +484,26 @@ def process_comment(
     return comment
 
 
+def function_def_end(source: str) -> int:
+    fd = ast.parse(textwrap.dedent(source)).body[0]
+    if not hasattr(fd, "args"):
+        return -1
+    return fd.args.args[-1].end_lineno if fd.args.args else -1
+
+
 def extract_code_params(code: Callable[..., Any]) -> list[Param]:
     ordered_params = code_to_ordered_params(code)
     hints = {k: v.annotation for k, v in ordered_params.items()}.copy()
     comment = ""
     param = None
     params: list[Param] = []
+    source = inspect.getsource(code)
+    fd_end = function_def_end(source)
 
-    for token in tokenize_string(inspect.getsource(code)):
+    for token in tokenize_string(source):
         if token.exact_type is COMMENT:
+            if fd_end > -1 and token.end[0] > fd_end:
+                break
             comment = process_comment(param, params, token)
             continue
         # tokenize.NL -
diff --git a/tests/test_extract_args.py b/tests/test_extract_args.py
index d5dad18..a40ec6b 100644
--- a/tests/test_extract_args.py
+++ b/tests/test_extract_args.py
@@ -458,3 +458,18 @@ def code(
             description="this is B",
         ),
      ]
+
+
+def test_two_args_unrelated_comment():
+    def code(
+        a: int,
+        b: int,
+    ):
+        # This is not an attribute comment
+        pass
+
+    params = simplecli.extract_code_params(code)
+    assert params == [
+        simplecli.Param(name="a", annotation=int),
+        simplecli.Param(name="b", annotation=int),
+     ]