flatten_tuple_return handles Return and Starred reader contexts

nominal_refactor_advisor/refactors/return_shape.py

@@ -281,6 +281,13 @@ def _rewrite_reader(
         if isinstance(parent, (ast.If, ast.BoolOp)):
             return
 
+        if isinstance(parent, ast.Return) and parent.value is call_node:
+            self._rewrite_return(use, parent, geometry, replacements)
+            return
+
+        if isinstance(parent, ast.Starred) and parent.value is call_node:
+            return
+
         raise ValueError(
             f"{use.where} reads `{self.method_name}` in a context this "
             f"operation cannot rewrite: {type(parent).__name__} at line "
@@ -333,4 +340,28 @@ def _rewrite_next_iter(
             )
         )
 
+    def _rewrite_return(
+        self,
+        use: AnsweredUse,
+        return_node: ast.Return,
+        geometry: SourceTextGeometry,
+        replacements: dict[str, list[SourceTextSpanReplacement]],
+    ) -> None:
+        """``return method()`` → ``return ((_r := method()),) if _r is not None else ()``
+
+        The caller still returns a tuple; the called method now returns
+        ``T | None``.  A walrus wraps the result back into a one-element
+        tuple or an empty tuple so the caller's contract stays unchanged."""
+        call_src = ast.unparse(use.node)
+        assert return_node.value is not None
+        start, end = geometry.required_node_offsets(return_node.value)
+        new_src = f"((_r := {call_src}),) if _r is not None else ()"
+        replacements.setdefault(use.file_path, []).append(
+            SourceTextSpanReplacement.from_offsets(
+                start_offset=start,
+                end_offset=end,
+                replacement_source=new_src,
+            )
+        )
+
 

tests/test_flatten_tuple_return.py

@@ -161,3 +161,29 @@ def test_annotation_flattens(tmp_path: Path) -> None:
     assert FlattenTupleReturnOperation._flatten_annotation("tuple[int, ...]") == "int | None"
     assert FlattenTupleReturnOperation._flatten_annotation("str") == "str"
     assert FlattenTupleReturnOperation._flatten_annotation("list[str]") == "list[str]"
+
+
+def test_return_context_reader_wraps_in_tuple(tmp_path: Path) -> None:
+    """A caller that does ``return method()`` wraps the flattened result back."""
+    package = tmp_path / "pkg"
+    package.mkdir()
+    (package / "__init__.py").write_text("")
+    (package / "base.py").write_text(BASE)
+    (package / "strict.py").write_text(STRICT)
+    (package / "consumer.py").write_text(
+        "from pkg.base import Validator\n\n\n"
+        "def relay(v: Validator) -> tuple[str, ...]:\n"
+        "    return v.warning()\n"
+    )
+    snapshot = CodemodSourceSnapshot.from_modules(parse_python_modules(tmp_path))
+    operation = _operation(tmp_path)
+    recipe = RefactorRecipe(recipe_id="test", operations=(operation,))
+    document = CodemodPlanDocument(recipes=(recipe,))
+    simulation = document.simulate(snapshot)
+    consumer_path = (tmp_path / "pkg/consumer.py").as_posix()
+    rewritten = simulation.simulation.rewritten_sources.get(
+        consumer_path, snapshot.sources_by_file_path[consumer_path]
+    )
+    assert "return v.warning()" not in rewritten
+    assert "_r := v.warning()" in rewritten
+    assert "is not None" in rewritten