Fixes #3 - Incorrect handling for first comment after function defini…

…tion (#13)

* Use ast to detect the last line of the last argument, stop iterating once we're past it

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 -

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),
+     ]