[Call-by-name] Fixed ignored compound statements and generator MultiGets (#21042)

This PR captures two missing cases from migration. 

- Generator-based `MultiGet`s (very important)
- `Get` calls within compound statements (helps with edge cases)

Author: sureshjoshi

src/python/pants/goal/migrate_call_by_name.py

@@ -12,7 +12,7 @@
 from dataclasses import dataclass
 from functools import partial
 from pathlib import Path, PurePath
-from typing import Callable, TypedDict
+from typing import Callable, Literal, TypedDict, cast
 
 from pants.base.build_environment import get_buildroot
 from pants.base.exiter import PANTS_SUCCEEDED_EXIT_CODE, ExitCode
@@ -225,8 +225,6 @@ def _perform_replacements_on_file(self, file: Path, replacements: list[Replaceme
         This function is not idempotent, so it should be run only once per file (per migration plan).
 
         Replacement imports are bulk added below the existing "pants.engine.rules" import.
-
-
         """
 
         imports_added = False
@@ -241,7 +239,7 @@ def _perform_replacements_on_file(self, file: Path, replacements: list[Replaceme
                 modified = False
                 for replacement in replacements:
                     if line_number == replacement.line_range[0]:
-                        line_end = ",\n" if replacement.is_argument else "\n"
+                        line_end = ",\n" if replacement.add_trailing_comma else "\n"
                         # On the first line of the range, emit the new source code where the old "Get" started
                         print(line[: replacement.col_range[0]], end="")
                         print(ast.unparse(replacement.new_source), end=line_end)
@@ -321,7 +319,8 @@ class Replacement:
     current_source: ast.Call
     new_source: ast.Call
     additional_imports: list[ast.ImportFrom]
-    is_argument: bool = False
+    # TODO: Use libcst or another CST, rather than an ast
+    add_trailing_comma: bool = False
 
     def sanitized_imports(self) -> list[ast.ImportFrom]:
         """Remove any circular or self-imports."""
@@ -361,7 +360,7 @@ def __str__(self) -> str:
             current_source={ast.dump(self.current_source, indent=2)},
             new_source={ast.dump(self.new_source, indent=2)},
             additional_imports={[ast.dump(i, indent=2) for i in self.additional_imports]},
-            is_argument={self.is_argument}
+            add_trailing_comma={self.add_trailing_comma}
         )
         """
 
@@ -613,19 +612,33 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):
         if not self._should_visit_node(node.decorator_list):
             return
 
-        for child in node.body:
-            if call := self._maybe_replaceable_call(child):
-                if replacement := self.syntax_mapper.map_get_to_new_syntax(
-                    call, self.filename, calling_func=node.name
-                ):
-                    self.replacements.append(replacement)
+        self._recurse_body_statements(node.name, node)
+
+    def _recurse_body_statements(self, root: str, node: ast.stmt):
+        """Recursively walk the body of a node, including properties of compound statements looking
+        for Get() calls to replace.
+
+        https://docs.python.org/3/reference/compound_stmts.html
+        """
+        for prop in ["body", "handlers", "orelse", "finalbody"]:
+            for child in getattr(node, prop, []):
+                self._recurse_body_statements(root, cast(ast.stmt, child))
 
-            for call in self._maybe_replaceable_multiget(child):
-                if replacement := self.syntax_mapper.map_get_to_new_syntax(
-                    call, self.filename, calling_func=node.name
-                ):
-                    replacement.is_argument = True
-                    self.replacements.append(replacement)
+        self._maybe_add_replacements(root, node)
+
+    def _maybe_add_replacements(self, calling_func: str, statement: ast.stmt):
+        if call := self._maybe_replaceable_call(statement):
+            if replacement := self.syntax_mapper.map_get_to_new_syntax(
+                call, self.filename, calling_func=calling_func
+            ):
+                self.replacements.append(replacement)
+
+        for call, statement_type in self._maybe_replaceable_multiget(statement):
+            if replacement := self.syntax_mapper.map_get_to_new_syntax(
+                call, self.filename, calling_func=calling_func
+            ):
+                replacement.add_trailing_comma = statement_type == "call"
+                self.replacements.append(replacement)
 
     def _should_visit_node(self, decorator_list: list[ast.expr]) -> bool:
         """Only interested in async functions with the @rule(...) or @goal_rule(...) decorator."""
@@ -667,23 +680,37 @@ def _maybe_replaceable_call(self, statement: ast.stmt) -> ast.Call | None:
             return call_node
         return None
 
-    def _maybe_replaceable_multiget(self, statement: ast.stmt) -> list[ast.Call]:
-        """Looks for the following form of MultiGet that we want to replace:
+    def _maybe_replaceable_multiget(
+        self, statement: ast.stmt
+    ) -> list[tuple[ast.Call, Literal["call", "generator"]]]:
+        """Looks for the following forms of MultiGet that we want to replace:
 
         - multigot = await MultiGet(Get(...), Get(...), ...)
+        - multigot = await MultiGet(Get(...) for x in y)
         """
-        if (
+        if not (
             isinstance(statement, ast.Assign)
             and isinstance((await_node := statement.value), ast.Await)
             and isinstance((call_node := await_node.value), ast.Call)
             and isinstance(call_node.func, ast.Name)
             and call_node.func.id == "MultiGet"
         ):
-            return [
-                arg
-                for arg in call_node.args
-                if isinstance(arg, ast.Call)
+            return []
+
+        args: list[tuple[ast.Call, Literal["call", "generator"]]] = []
+        for arg in call_node.args:
+            if (
+                isinstance(arg, ast.Call)
                 and isinstance(arg.func, ast.Name)
                 and arg.func.id == "Get"
-            ]
-        return []
+            ):
+                args.append((arg, "call"))
+
+            if (
+                isinstance(arg, ast.GeneratorExp)
+                and isinstance((call_arg := arg.elt), ast.Call)
+                and isinstance(call_arg.func, ast.Name)
+                and call_arg.func.id == "Get"
+            ):
+                args.append((call_arg, "generator"))
+        return args

src/python/pants/goal/migrate_call_by_name_integration_test.py

@@ -95,6 +95,20 @@ async def variants(black: Black, local_env: ChosenLocalEnvironmentName) -> Foo:
         pex = await Get(VenvPex, PexRequest, black.to_pex_request())
         digest = await Get(Digest, CreateArchive(EMPTY_SNAPSHOT))
         paths = await Get(BinaryPaths, {{BinaryPathRequest(binary_name="time", search_path=("/usr/bin")): BinaryPathRequest, local_env.val: EnvironmentName}})
+        try:
+            try_all_targets_try = await Get(AllTargets)
+        except:
+            try_all_targets_except = await Get(AllTargets)
+        else:
+            try_all_targets_else = await Get(AllTargets)
+        finally:
+            try_all_targets_finally = await Get(AllTargets)
+        if True:
+            conditional_all_targets_if = await Get(AllTargets)
+        elif False:
+            conditional_all_targets_elif = await Get(AllTargets)
+        else:
+            conditional_all_targets_else = await Get(AllTargets)
 
     class Bar:
         pass
@@ -136,6 +150,10 @@ async def multiget(black: Black) -> Thud:
             ),
             digest_get
         )
+        multigot_forloop = await MultiGet(
+            Get(Digest, CreateArchive(EMPTY_SNAPSHOT))
+            for i in [0, 1, 2]
+        )
 
     def rules():
         return collect_rules()
@@ -168,6 +186,20 @@ async def variants(black: Black, local_env: ChosenLocalEnvironmentName) -> Foo:
         pex = await create_venv_pex(**implicitly({black.to_pex_request(): PexRequest}))
         digest = await create_archive(CreateArchive(EMPTY_SNAPSHOT), **implicitly())
         paths = await find_binary(**implicitly({BinaryPathRequest(binary_name='time', search_path='/usr/bin'): BinaryPathRequest, local_env.val: EnvironmentName}))
+        try:
+            try_all_targets_try = await find_all_targets(**implicitly())
+        except:
+            try_all_targets_except = await find_all_targets(**implicitly())
+        else:
+            try_all_targets_else = await find_all_targets(**implicitly())
+        finally:
+            try_all_targets_finally = await find_all_targets(**implicitly())
+        if True:
+            conditional_all_targets_if = await find_all_targets(**implicitly())
+        elif False:
+            conditional_all_targets_elif = await find_all_targets(**implicitly())
+        else:
+            conditional_all_targets_else = await find_all_targets(**implicitly())
 
     class Bar:
         pass
@@ -201,6 +233,10 @@ async def multiget(black: Black) -> Thud:
             create_venv_pex(**implicitly({black.to_pex_request(): PexRequest})),
             digest_get
         )
+        multigot_forloop = await concurrently(
+            create_archive(CreateArchive(EMPTY_SNAPSHOT), **implicitly())
+            for i in [0, 1, 2]
+        )
 
     def rules():
         return collect_rules()
@@ -275,8 +311,9 @@ def test_migrate_call_by_name_syntax():
         # Ensure the JSON output contains the paths to the files we expect to migrate
         assert all(str(p) in result.stdout for p in [register_path, rules1_path, rules2_path])
         # Ensure the warning for embedded comments is logged
+        # Note: This assertion is brittle - adding extra content to rules1.py will probably mean updating the range
         assert (
-            f"Comments found in {tmpdir}/src/migrateme/rules1.py within replacement range: (37, 42)"
+            f"Comments found in {tmpdir}/src/migrateme/rules1.py within replacement range: (51, 56)"
             in result.stderr
         )