From dc9fe89938c929818eac39d4b8f5cc6a9a3d95a6 Mon Sep 17 00:00:00 2001
From: Claude <noreply@anthropic.com>
Date: Wed, 11 Nov 2026 19:20:55 +0000
Subject: [PATCH] flatten_tuple_return handles Return and Starred reader
 contexts

_rewrite_reader now handles two more patterns:
- Return: `return method()` wraps the flattened result back into a
  tuple using a walrus: `((_r := method()),) if _r is not None else ()`
- Starred: `*method()` in a tuple is a passthrough (no rewrite needed)

These unblock most of the 13 zero_one_tuple findings on agent-comms
where callers use `return method()` (e.g. RuntimeReference.hidden at
runtime_processes.py:310) or tuple spread.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01RnH6dR7JdknxC8XJJtW8Uh
---
 .../refactors/return_shape.py                 | 31 +++++++++++++++++++
 tests/test_flatten_tuple_return.py            | 26 ++++++++++++++++
 2 files changed, 57 insertions(+)

diff --git a/nominal_refactor_advisor/refactors/return_shape.py b/nominal_refactor_advisor/refactors/return_shape.py
index 5944cb28..12ca83b6 100644
--- a/nominal_refactor_advisor/refactors/return_shape.py
+++ b/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,
+            )
+        )
+
 
diff --git a/tests/test_flatten_tuple_return.py b/tests/test_flatten_tuple_return.py
index 9719ceda..54171e5a 100644
--- a/tests/test_flatten_tuple_return.py
+++ b/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