fix: JS/TS augmentor robustness (export default class, duplicate methods, generators, class-field arrows, default params) + Python grouped imports

Author: Shtirmann

osa_tool/operations/codebase/docstring_generation/adapters/python_adapter.py

@@ -93,7 +93,13 @@ def _resolve_import_path(import_text, cwd):
             if not os.path.exists(module_path):
                 return import_mapping
 
-            for entity in (e.strip() for e in import_part.split(",")):
+            # a grouped import `from x import (a, b)` (possibly multi-line) keeps the
+            # parentheses/newlines in the raw text; strip them so the names stay clean
+            import_part = import_part.strip().strip("()").replace("\\", " ")
+
+            for entity in (e.strip().strip("()") for e in import_part.split(",")):
+                if not entity or entity == "*":
+                    continue
                 if " as " in entity:
                     imported_name, alias_name = (e.strip() for e in entity.split(" as ", 1))
                 else:

osa_tool/operations/codebase/docstring_generation/adapters/typescript_adapter.py

@@ -28,10 +28,15 @@ def get_name(self, node, sv):
         if n:
             return sv.text(n)
 
-        # arrow functions / function expressions have no name field;
-        # the name usually sits on the enclosing declarator (const foo = () => ...)
+        # arrow functions / function expressions have no name field; the name usually
+        # sits on the enclosing declarator (const foo = () => ...) or, for a class field
+        # arrow (foo = () => ...), on the surrounding field definition.
         parent = node.parent
-        if parent and parent.type == "variable_declarator":
+        if parent and parent.type in (
+            "variable_declarator",
+            "public_field_definition",
+            "field_definition",
+        ):
             declared = parent.child_by_field_name("name")
             if declared:
                 return sv.text(declared)

osa_tool/operations/codebase/docstring_generation/insert/ts_js_augmentor.py

@@ -23,7 +23,7 @@ def augment(self, file: str, source_code: str, docstrings: dict) -> dict[str, st
     def _inject_classes(self, lines, classes):
         for doc, class_name in classes:
 
-            class_pattern = re.compile(rf"^\s*(export\s+)?(abstract\s+)?class\s+{re.escape(class_name)}\b")
+            class_pattern = re.compile(rf"^\s*(export\s+)?(default\s+)?(abstract\s+)?class\s+{re.escape(class_name)}\b")
 
             for i, line in enumerate(lines):
                 if class_pattern.search(line):
@@ -59,6 +59,11 @@ def _inject_functions(self, lines, functions):
         return lines
 
     def _inject_methods(self, lines, methods):
+        # Decide targets first (pure scan, no mutation) so `used` line indices stay
+        # stable, then apply bottom-up so earlier indices remain valid while we edit.
+        used = set()
+        actions = []
+
         for doc, meta in methods:
             name = meta["method_name"]
             patterns = [
@@ -89,6 +94,19 @@ def _inject_methods(self, lines, methods):
                     """,
                     re.VERBOSE,
                 ),
+                # generator method: * method<T>(  /  async * method(  /  static * method(
+                re.compile(
+                    rf"""^\s*
+                    (?:public|private|protected|static|readonly|async|\s)*
+                    \*\s*
+                    {re.escape(name)}
+                    \s*
+                    (?:<[^>]*>)?
+                    \s*
+                    \(
+                    """,
+                    re.VERBOSE,
+                ),
                 # method<T>(
                 re.compile(
                     rf"""^\s*
@@ -111,10 +129,21 @@ def _inject_methods(self, lines, methods):
                     re.VERBOSE,
                 ),
             ]
+            # class field arrow: name = (...) =>  /  name = async (...) =>  /  name = x =>
+            field_arrow = re.compile(
+                rf"^\s*(?:(?:public|private|protected|static|readonly)\s+)*"
+                rf"{re.escape(name)}\s*=\s*(?:async\s+)?"
+                rf"(?:\([^)]*\)|[A-Za-z_$][\w$]*)\s*(?::\s*[^=]+?)?\s*=>"
+            )
 
             for i, line in enumerate(lines):
+                if i in used:
+                    continue
+
                 stripped = line.strip()
-                # skip obvious calls/usages
+                is_field_arrow = bool(field_arrow.search(line))
+
+                # skip obvious calls / control-flow / usages
                 if (
                     stripped.startswith("return ")
                     or stripped.startswith("if ")
@@ -124,17 +153,31 @@ def _inject_methods(self, lines, methods):
                     or stripped.startswith("catch ")
                     or stripped.startswith("new ")
                     or re.search(rf"\.\s*{re.escape(name)}\s*\(", stripped)
-                    or "=" in stripped
                 ):
                     continue
 
-                if any(p.search(line) for p in patterns):
-                    if self._has_doc(lines, i):
-                        lines = self._replace_doc(lines, i, doc)
-                    else:
-                        lines.insert(i, self._format(doc, line))
+                # a statement (ends with ';') or an assignment ('=' before the call
+                # parens) is a call/usage, not a declaration -- skip it. A legitimate
+                # class-field arrow is allowed through (it also ends with ';').
+                if not is_field_arrow:
+                    if stripped.endswith(";"):
+                        continue
+                    paren = stripped.find("(")
+                    before_paren = stripped if paren == -1 else stripped[:paren]
+                    if "=" in before_paren:
+                        continue
+
+                if is_field_arrow or any(p.search(line) for p in patterns):
+                    used.add(i)
+                    actions.append((i, doc))
                     break
 
+        for i, doc in sorted(actions, key=lambda a: a[0], reverse=True):
+            if self._has_doc(lines, i):
+                lines = self._replace_doc(lines, i, doc)
+            else:
+                lines.insert(i, self._format(doc, lines[i]))
+
         return lines
 
     def _has_doc(self, lines, i):

tests/unit/operations/codebase/docstring_generation/test_adapters_robustness.py

@@ -0,0 +1,64 @@
+from osa_tool.operations.codebase.docstring_generation.adapters.python_adapter import PythonAdapter
+from osa_tool.operations.codebase.docstring_generation.core.osa_parser import OSA_TreeSitter
+
+# --- #5: class-field arrow name resolution (TypeScript adapter) --------------
+
+
+def _method_names(res):
+    names = []
+    for item in res["structure"]:
+        if item.get("type") == "class":
+            names += [m["method_name"] for m in item["methods"]]
+        elif item.get("type") == "function":
+            names.append(item["details"]["method_name"])
+    return names
+
+
+def test_class_field_arrow_gets_name(tmp_path):
+    f = tmp_path / "c.ts"
+    f.write_text(
+        "class C {\n  greet = (name: string): string => 'hi ' + name;\n}\n",
+        encoding="utf-8",
+    )
+    res = OSA_TreeSitter(str(tmp_path)).extract_structure(str(f))
+    names = _method_names(res)
+    assert "greet" in names
+    assert "anonymous" not in names
+
+
+def test_regular_arrow_const_still_named(tmp_path):
+    f = tmp_path / "u.ts"
+    f.write_text("export const mul = (a, b) => a * b;\n", encoding="utf-8")
+    res = OSA_TreeSitter(str(tmp_path)).extract_structure(str(f))
+    assert "mul" in _method_names(res)
+    assert "anonymous" not in _method_names(res)
+
+
+# --- #6: parenthesized / multi-line Python from-imports ----------------------
+
+
+def test_grouped_import_clean_keys(tmp_path):
+    (tmp_path / "mod.py").write_text("a = 1\nb = 2\n", encoding="utf-8")
+    mapping = PythonAdapter._resolve_import_path("from mod import (a, b)", str(tmp_path))
+    assert set(mapping.keys()) == {"a", "b"}
+    assert not any("(" in k or ")" in k for k in mapping)
+
+
+def test_multiline_grouped_import_clean_keys(tmp_path):
+    (tmp_path / "mod.py").write_text("a = 1\nb = 2\nc = 3\n", encoding="utf-8")
+    text = "from mod import (\n    a,\n    b,\n    c,\n)"
+    mapping = PythonAdapter._resolve_import_path(text, str(tmp_path))
+    assert set(mapping.keys()) == {"a", "b", "c"}
+
+
+def test_grouped_import_with_alias(tmp_path):
+    (tmp_path / "mod.py").write_text("a = 1\nb = 2\n", encoding="utf-8")
+    mapping = PythonAdapter._resolve_import_path("from mod import (a as x, b)", str(tmp_path))
+    assert set(mapping.keys()) == {"x", "b"}
+    assert mapping["x"]["class"] == "a"
+
+
+def test_plain_import_still_works(tmp_path):
+    (tmp_path / "mod.py").write_text("a = 1\n", encoding="utf-8")
+    mapping = PythonAdapter._resolve_import_path("from mod import a", str(tmp_path))
+    assert set(mapping.keys()) == {"a"}

tests/unit/operations/codebase/docstring_generation/test_ts_js_augmentor.py

@@ -112,3 +112,124 @@ def test_format_plain_text_unchanged_behaviour():
     assert "```" not in out
     assert out.count("/**") == 1
     assert out.count("*/") == 1
+
+
+# --- robustness fixes: class/method injection edge cases ---------------------
+
+
+def _line_above(out, decl_substring):
+    """Return the stripped line directly above the first line containing decl_substring."""
+    lines = out.splitlines()
+    for i, l in enumerate(lines):
+        if decl_substring in l:
+            return lines[i - 1].strip() if i > 0 else ""
+    return None
+
+
+def test_inject_export_default_class():
+    """`export default class X` must receive a class-level JSDoc (regex #1)."""
+    aug = TSJSAugmentor()
+    src = "export default class Queue {\n\tclear() {}\n}\n"
+    out = aug.augment("f.ts", src, {"classes": [("A FIFO queue.", "Queue")]})["f.ts"]
+
+    assert "A FIFO queue." in out
+    assert _line_above(out, "export default class Queue") == "*/"
+
+
+def test_inject_generator_method():
+    """Generator methods (`* gen()`) must be matched and documented (#4)."""
+    aug = TSJSAugmentor()
+    src = "class Q {\n\t* drain() {\n\t\tyield 1;\n\t}\n}\n"
+    out = aug.augment("f.ts", src, {"methods": [("Drains the queue.", {"method_name": "drain"})]})["f.ts"]
+
+    assert "Drains the queue." in out
+    assert _line_above(out, "* drain()") == "*/"
+
+
+def test_inject_two_same_named_methods_both_documented():
+    """Two methods with the same name must each get their own doc, in order (#2)."""
+    aug = TSJSAugmentor()
+    src = (
+        "class Node {\n\tconstructor(value) {\n\t\tthis.value = value;\n\t}\n}\n"
+        "class Queue {\n\tconstructor() {\n\t\tthis.clear();\n\t}\n}\n"
+    )
+    docs = {
+        "methods": [
+            ("Node constructor.", {"method_name": "constructor"}),
+            ("Queue constructor.", {"method_name": "constructor"}),
+        ]
+    }
+    out = aug.augment("f.ts", src, docs)["f.ts"]
+
+    assert out.count("constructor(") == 2
+    assert "Node constructor." in out
+    assert "Queue constructor." in out
+    # both declarations carry a doc block directly above them
+    lines = out.splitlines()
+    ctor_idxs = [i for i, l in enumerate(lines) if "constructor(" in l]
+    assert len(ctor_idxs) == 2
+    for idx in ctor_idxs:
+        assert lines[idx - 1].strip() == "*/"
+    # order preserved: first constructor -> first doc
+    assert out.index("Node constructor.") < out.index("Queue constructor.")
+
+
+def test_inject_method_with_default_param_value():
+    """A method whose signature line contains `=` (default value) must be documented (#3)."""
+    aug = TSJSAugmentor()
+    src = "class C {\n\tscale(x, factor = 2) {\n\t\treturn x * factor;\n\t}\n}\n"
+    out = aug.augment("f.ts", src, {"methods": [("Scales a value.", {"method_name": "scale"})]})["f.ts"]
+
+    assert "Scales a value." in out
+    assert _line_above(out, "scale(x, factor = 2)") == "*/"
+
+
+def test_assignment_is_not_mistaken_for_method():
+    """A local assignment must NOT be documented as a method (guard still holds)."""
+    aug = TSJSAugmentor()
+    src = "class C {\n\trun() {\n\t\tconst compute = helper(1);\n\t\treturn compute;\n\t}\n}\n"
+    out = aug.augment("f.ts", src, {"methods": [("Computes.", {"method_name": "compute"})]})["f.ts"]
+
+    assert "Computes." not in out
+
+
+def test_inject_class_field_arrow_method():
+    """A class-field arrow (`name = (...) => ...`) must be documented (#5, injection)."""
+    aug = TSJSAugmentor()
+    src = "class C {\n\tgreet = (name) => 'hi ' + name;\n}\n"
+    out = aug.augment("f.ts", src, {"methods": [("Greets by name.", {"method_name": "greet"})]})["f.ts"]
+
+    assert "Greets by name." in out
+    assert _line_above(out, "greet = (name) =>") == "*/"
+
+
+def test_bare_call_with_operator_not_documented():
+    """A bare call statement whose args contain `===`/`>=`/etc must NOT be mistaken
+    for a method declaration; the real same-named method gets the doc instead (#3 guard)."""
+    aug = TSJSAugmentor()
+    src = (
+        "class S {\n"
+        "\tboot() {\n"
+        "\t\tassert(ready === true);\n"
+        "\t}\n"
+        "\tassert(cond) {\n"
+        "\t\treturn cond;\n"
+        "\t}\n"
+        "}\n"
+    )
+    docs = {
+        "methods": [
+            ("Boots the service.", {"method_name": "boot"}),
+            ("Asserts a condition.", {"method_name": "assert"}),
+        ]
+    }
+    out = aug.augment("f.ts", src, docs)["f.ts"]
+    lines = out.splitlines()
+
+    # the call statement must NOT receive a doc
+    call_idx = next(i for i, l in enumerate(lines) if "assert(ready === true)" in l)
+    assert lines[call_idx - 1].strip() != "*/"
+    # the real declaration does
+    decl_idx = next(i for i, l in enumerate(lines) if "assert(cond)" in l)
+    assert lines[decl_idx - 1].strip() == "*/"
+    assert "Asserts a condition." in out