19381742489

Committed 15 Nov 2025 12:52AM UTC coverage: 49.706% (-30.6%) from 80.29%

Build # 19381742489

Build Type

Pull #22890

github

Committed by

web-flow

Commit Message

Merge d961abf79 into 42e1ebd41

Pull Request Pull Request #22890: Updated all python subsystem constraints to 3.14

Run Details

4 of 5 new or added lines in 5 files covered. (80.0%)

14659 existing lines in 485 files now uncovered.

31583 of 63540 relevant lines covered (49.71%)

0.79 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

28.41

/src/python/pants/goal/migrate_call_by_name.py

# Copyright 2024 Pants project contributors (see CONTRIBUTORS.md).
# Licensed under the Apache License, Version 2.0 (see LICENSE).

from __future__ import annotations

import importlib.util
import json
import logging
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from functools import partial
from pathlib import Path, PurePath
from typing import TypedDict

import libcst as cst
import libcst.helpers as h
import libcst.matchers as m
import libcst.metadata
from libcst.display import dump

from pants.base.build_environment import get_buildroot
from pants.base.exiter import PANTS_SUCCEEDED_EXIT_CODE, ExitCode
from pants.base.specs import Specs
from pants.build_graph.build_configuration import BuildConfiguration
from pants.engine.fs import Paths
from pants.engine.unions import UnionMembership
from pants.goal.builtin_goal import BuiltinGoal
from pants.init.engine_initializer import GraphSession
from pants.option.option_types import BoolOption
from pants.option.options import Options
from pants.util import cstutil
from pants.util.strutil import softwrap

logger = logging.getLogger(__name__)


class MigrateCallByNameBuiltinGoal(BuiltinGoal):
    name = "migrate-call-by-name"
    help = softwrap(
        """
        Migrate from `Get` syntax to call-by-name syntax (#19730). This is a **destructive** operation,
        so only run this on source controlled files that you are prepared to revert if necessary.

        This goal will attempt to migrate the set of paths/targets specified at the command line
        if they are part of the "migration plan". This migration does not add any new files, but
        instead modifies existing files in-place without any formatting. The resulting changes should
        be reviewed, tested, and formatted/linted before committing.

        The migration plan is a JSON representation of the rule graph, which is generated by the
        engine based on the active backends/rules in the project.

        Each item in the migration plan is a rule that contains the old `Get` syntax, the associated
        input/output types, and the new function to directly call. The migration plan can be dumped as
        JSON using the `--json` flag, which can be useful for debugging. For example:

        {
            "filepath": "src/python/pants/source/source_root.py",
            "function": "get_source_roots",
            "gets": [{
                "input_types": [{ "module": "pants.source.source_root", "name": "SourceRootsRequest" }],
                "output_type": { "module": "pants.source.source_root", "name": "OptionalSourceRootsResult" },
                "rule_dep": { "function": "get_optional_source_roots", "module": "pants.source.source_root" }
            }],
            "module": "pants.source.source_root"
        }
        """
    )

    should_dump_json = BoolOption(
        flag_name="--json", help=softwrap("Dump the migration plan as JSON"), default=False
    )

    def run(
        self,
        *,
        build_config: BuildConfiguration,
        graph_session: GraphSession,
        options: Options,
        specs: Specs,
        union_membership: UnionMembership,
    ) -> ExitCode:
        migration_plan = self._create_migration_plan(graph_session, PurePath(get_buildroot()))
        if self.should_dump_json:
            print(json.dumps(migration_plan, indent=2, sort_keys=True))

        path_globs = specs.includes.to_specs_paths_path_globs()
        if not path_globs.globs:
            return PANTS_SUCCEEDED_EXIT_CODE

        plan_files = {item["filepath"] for item in migration_plan}

        paths: list[Paths] = graph_session.scheduler_session.product_request(Paths, path_globs)
        requested_files = set(paths[0].files)

        files_to_migrate = requested_files.intersection(plan_files)
        if not files_to_migrate:
            logger.info(
                f"None of the {len(requested_files)} requested files are part of the {len(plan_files)} files in the migration plan. Please ensure the backend containing these files is activated in pants.toml."
            )
            return PANTS_SUCCEEDED_EXIT_CODE

        syntax_mapper = CallByNameSyntaxMapper(migration_plan)
        for f in sorted(files_to_migrate):
            file = Path(f)
            logger.info(f"Processing {file}")

            transformer = CallByNameTransformer(file, syntax_mapper)
            source_code = Path.read_text(file)
            source_tree = cst.parse_module(source_code)
            new_tree = source_tree.visit(transformer)

            if not new_tree.deep_equals(source_tree):
                new_source = new_tree.code
                Path.write_text(file, new_source)

        return PANTS_SUCCEEDED_EXIT_CODE

    def _create_migration_plan(
        self, session: GraphSession, build_root: PurePath
    ) -> list[RuleGraphGet]:
        """Use the rule graph to create a migration plan for each "active" file that uses the old
        Get() syntax.

        This function is mostly about creating a stable-sorted collection of items with metadata for
        downstream
        """
        items: list[RuleGraphGet] = []
        for rule, deps in session.scheduler_session.rule_graph_rule_gets().items():
            if isinstance(rule, partial):
                # Ignoring partials, see https://github.com/pantsbuild/pants/issues/20744
                continue

            assert (spec := importlib.util.find_spec(rule.__module__)) is not None
            assert spec.origin is not None

            try:
                spec_origin = str(PurePath(spec.origin).relative_to(build_root))
            except ValueError:
                logger.debug(
                    f"Ignoring migration plan item located outside of build_root ({build_root}) - file was located at {spec.origin}"
                )
                continue

            item: RuleGraphGet = {
                "filepath": spec_origin,
                "module": rule.__module__,
                "function": rule.__name__,
                "gets": [],
            }
            unsorted_deps: list[RuleGraphGetDep] = []

            for output_type, input_types, rule_dep in deps:
                if isinstance(rule_dep, partial):
                    # Ignoring partials, see https://github.com/pantsbuild/pants/issues/20744
                    continue

                unsorted_deps.append(
                    {
                        "input_types": sorted(
                            [
                                {
                                    "module": input_type.__module__,
                                    "name": input_type.__name__,
                                }
                                for input_type in input_types
                            ],
                            key=lambda x: (x["module"], x["name"]),
                        ),
                        "output_type": {
                            "module": output_type.__module__,
                            "name": output_type.__name__,
                        },
                        "rule_dep": {
                            "function": rule_dep.__name__,
                            "module": rule_dep.__module__,
                        },
                    }
                )

            sorted_deps = sorted(
                unsorted_deps, key=lambda x: (x["rule_dep"]["module"], x["rule_dep"]["function"])
            )
            item["gets"] = sorted_deps
            items.append(item)

        return sorted(items, key=lambda x: (x["filepath"], x["function"]))


# ------------------------------------------------------------------------------------------
# Migration Plan Typed Dicts
# ------------------------------------------------------------------------------------------


class RuleGraphGet(TypedDict):
    filepath: str
    function: str
    module: str
    gets: list[RuleGraphGetDep]


class RuleGraphGetDep(TypedDict):
    input_types: list[RuleType]
    output_type: RuleType
    rule_dep: RuleFunction


class RuleType(TypedDict):
    module: str
    name: str


class RuleFunction(TypedDict):
    function: str
    module: str


# ------------------------------------------------------------------------------------------
# Replacement container
# ------------------------------------------------------------------------------------------


@dataclass
class Replacement:
    filename: PurePath
    module: str
    current_source: cst.Call
    new_source: cst.Call
    additional_imports: list[cst.ImportFrom]

    def sanitized_imports(self) -> list[cst.ImportFrom]:
        """Remove any circular or self-imports."""
        cst_module = cstutil.make_importfrom_attr(self.module)
        return [
            i for i in self.additional_imports if i.module and not cst_module.deep_equals(i.module)
        ]

    def sanitize(self, unavailable_names: set[str]):
        """Remove any shadowing of names, except if the new_func is in the current file."""

        func_name = cst.ensure_type(self.new_source.func, cst.Name)
        if func_name.value not in unavailable_names:
            return

        # If the new func_name is not in the sanitized imports, it must already be in the current file
        imported_names: set[str] = set()
        for imp in self.sanitized_imports():
            assert isinstance(imp.names, Iterable)
            for import_alias in imp.names:
                alias_name = cst.ensure_type(import_alias.name, cst.Name)
                imported_names.add(alias_name.value)

        if func_name.value not in imported_names:
            return

        # In-place update this replacement and additional imports
        bound_name = f"{func_name.value}_get"
        self.new_source = self.new_source.with_deep_changes(self.new_source.func, value=bound_name)

        for i, imp in enumerate(self.additional_imports):
            assert isinstance(imp.names, Iterable)
            for import_alias in imp.names:
                alias_name = cst.ensure_type(import_alias.name, cst.Name)
                if alias_name.value == func_name.value:
                    self.additional_imports[i] = imp.with_changes(
                        names=[cst.ImportAlias(func_name, asname=cst.AsName(cst.Name(bound_name)))]
                    )

        logging.warning(f"Renamed {func_name} to {bound_name} to avoid shadowing")

    def __str__(self) -> str:
        return f"""
        Replacement(
            filename={self.filename},
            module={self.module},
            current_source={dump(self.current_source)},
            new_source={dump(self.new_source)},
            additional_imports={[dump(i) for i in self.additional_imports]},
        )
        """


# ------------------------------------------------------------------------------------------
# Call-by-name transformer
# ------------------------------------------------------------------------------------------


class CallByNameTransformer(m.MatcherDecoratableTransformer):
    def __init__(self, filename: PurePath, syntax_mapper: CallByNameSyntaxMapper) -> None:
        super().__init__()

        self.filename = filename
        self.syntax_mapper = syntax_mapper
        self.calling_function: str = ""
        self.additional_imports: list[cst.ImportFrom] = []
        self.unavailable_names: set[str] = set()

    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
        self.calling_function = node.name.value

    def leave_FunctionDef(
        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
    ) -> cst.FunctionDef:
        self.calling_function = ""
        return updated_node

    @m.leave(m.Call(func=m.Name("Get")))
    def handle_get(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
        replacement = self.syntax_mapper.map_get_to_new_syntax(
            original_node, self.filename, self.calling_function
        )
        if not replacement:
            return updated_node

        replacement.sanitize(self.unavailable_names)
        self.additional_imports.extend(replacement.sanitized_imports())
        return replacement.new_source

    @m.leave(m.Name("MultiGet"))
    def handle_multiget(self, original_node: cst.Name, updated_node: cst.Name) -> cst.Name:
        return updated_node.with_changes(value="concurrently")

    def visit_Module(self, node: cst.Module):
        """Collects all names we risk shadowing."""
        wrapper = libcst.metadata.MetadataWrapper(module=node)
        scopes = set(wrapper.resolve(libcst.metadata.ScopeProvider).values())
        self.unavailable_names = {
            assignment.name for scope in scopes if scope for assignment in scope.assignments
        }

    def leave_Module(self, original_node: cst.Module, updated_node: cst.Module) -> cst.Module:
        """Performs final updates on imports, and sanitization."""
        if not self.additional_imports:
            return updated_node

        rules_import_index = 1
        for i, statement in enumerate(updated_node.body):
            if m.matches(
                statement,
                matcher=m.SimpleStatementLine(
                    body=[
                        m.ImportFrom(
                            module=cstutil.make_importfrom_attr_matcher("pants.engine.rules")
                        )
                    ]
                ),
            ):
                rules_import_index = i + 1
                break

        self.additional_imports.append(cstutil.make_importfrom("pants.engine.rules", "implicitly"))
        additional_import_statements = [
            cst.SimpleStatementLine(body=[i]) for i in self.additional_imports
        ]
        additional_import_statements = [
            v1
            for i, v1 in enumerate(additional_import_statements)
            if not any(v1.deep_equals(v2) for v2 in additional_import_statements[:i])
        ]

        return updated_node.with_changes(
            body=[
                *updated_node.body[:rules_import_index],
                *additional_import_statements,
                *updated_node.body[rules_import_index:],
            ]
        )


# ------------------------------------------------------------------------------------------
# Call-by-name syntax mapping
# ------------------------------------------------------------------------------------------


class CallByNameSyntaxMapper:
    def __init__(self, graphs: list[RuleGraphGet]) -> None:
        self.graphs = graphs

        self.mapping: dict[
            tuple[int, type[cst.Call] | type[cst.Dict] | None],
            Callable[[cst.Call, list[RuleGraphGetDep]], tuple[cst.Call, list[cst.ImportFrom]]],
        ] = {
            (1, None): self.map_no_args_get_to_new_syntax,
            (2, cst.Call): self.map_short_form_get_to_new_syntax,
            (2, cst.Dict): self.map_dict_form_get_to_new_syntax,
            (3, None): self.map_long_form_get_to_new_syntax,
        }

    def _get_graph_item(self, filename: PurePath, calling_func: str) -> RuleGraphGet | None:
        return next(
            (
                item
                for item in self.graphs
                if item["filepath"] == str(filename) and item["function"] == calling_func
            ),
            None,
        )

    def map_get_to_new_syntax(
        self, get: cst.Call, filename: PurePath, calling_func: str
    ) -> Replacement | None:
        new_source: cst.Call | None = None
        imports: list[cst.ImportFrom] = []

        if not (graph_item := self._get_graph_item(filename, calling_func)):
            logger.warning(f"Failed to find dependencies for {calling_func} in {filename}")
            return None

        get_deps = graph_item["gets"]
        num_args = len(get.args)

        arg_type = None
        if num_args == 2 and (arg_type := type(get.args[1].value)) not in [cst.Call, cst.Dict]:
            logger.warning(f"Failed to migrate: Unknown arg type {get.args[1]}")
            return None

        try:
            new_source, imports = self.mapping[(num_args, arg_type)](get, get_deps)  # type: ignore
        except Exception as e:
            logger.warning(
                f"Failed to migrate Get ({num_args}, {arg_type}) in {filename}:{calling_func} due to: {e}\n"
            )
            return None

        return Replacement(
            filename=filename,
            module=graph_item["module"],
            current_source=get,
            new_source=new_source,
            additional_imports=imports,
        )

    def map_no_args_get_to_new_syntax(
        self, get: cst.Call, deps: list[RuleGraphGetDep]
    ) -> tuple[cst.Call, list[cst.ImportFrom]]:
        """Get(<OutputType>) => the_rule_to_call(**implicitly())"""

        output_type = cst.ensure_type(get.args[0].value, cst.Name).value

        dep = next(
            dep
            for dep in deps
            if dep["output_type"]["name"] == output_type and not dep["input_types"]
        )
        module, new_function = dep["rule_dep"]["module"], dep["rule_dep"]["function"]

        new_call = cst.Call(
            func=cst.Name(new_function),
            args=[cst.Arg(value=cst.Call(cst.Name("implicitly")), star="**")],
        )
        if called_funcdef := cstutil.extract_functiondef_from_module(module, new_function):
            new_call = fix_implicitly_usage(new_call, called_funcdef)

        imports = [cstutil.make_importfrom(module, new_function)]
        return new_call, imports

    def map_long_form_get_to_new_syntax(
        self, get: cst.Call, deps: list[RuleGraphGetDep]
    ) -> tuple[cst.Call, list[cst.ImportFrom]]:
        """Get(<OutputType>, <InputType>, input) => the_rule_to_call(**implicitly(input))"""

        output_type = cst.ensure_type(get.args[0].value, cst.Name).value
        input_type = cst.ensure_type(get.args[1].value, cst.Name).value

        dep = next(
            dep
            for dep in deps
            if dep["output_type"]["name"] == output_type
            and len(dep["input_types"]) == 1
            and dep["input_types"][0]["name"] == input_type
        )
        module, new_function = dep["rule_dep"]["module"], dep["rule_dep"]["function"]

        new_call = cst.Call(
            func=cst.Name(new_function),
            args=[
                cst.Arg(
                    value=cst.Call(
                        func=cst.Name("implicitly"),
                        args=[
                            cst.Arg(
                                value=cst.Dict(
                                    [
                                        cst.DictElement(
                                            key=get.args[2].value, value=cst.Name(input_type)
                                        )
                                    ]
                                )
                            )
                        ],
                    ),
                    star="**",
                )
            ],
        )

        if called_funcdef := cstutil.extract_functiondef_from_module(module, new_function):
            new_call = fix_implicitly_usage(new_call, called_funcdef)

        imports = [cstutil.make_importfrom(module, new_function)]
        return new_call, imports

    def map_short_form_get_to_new_syntax(
        self, get: cst.Call, deps: list[RuleGraphGetDep]
    ) -> tuple[cst.Call, list[cst.ImportFrom]]:
        """Get(<OutType>, <InputType>(<input args>)) => the_rule_to_call(input, **implicitly())"""

        output_type = cst.ensure_type(get.args[0].value, cst.Name).value
        input_call = cst.ensure_type(get.args[1].value, cst.Call)
        input_type = cst.ensure_type(input_call.func, cst.Name).value

        dep = next(
            dep
            for dep in deps
            if dep["output_type"]["name"] == output_type
            and len(dep["input_types"]) == 1
            and dep["input_types"][0]["name"] == input_type
        )
        module, new_function = dep["rule_dep"]["module"], dep["rule_dep"]["function"]

        new_call = cst.Call(
            func=cst.Name(new_function),
            args=[
                cst.Arg(value=input_call),
                cst.Arg(value=cst.Call(cst.Name("implicitly")), star="**"),
            ],
        )

        if called_funcdef := cstutil.extract_functiondef_from_module(module, new_function):
            new_call = fix_implicitly_usage(new_call, called_funcdef)

        imports = [cstutil.make_importfrom(module, new_function)]
        return new_call, imports

    def map_dict_form_get_to_new_syntax(
        self, get: cst.Call, deps: list[RuleGraphGetDep]
    ) -> tuple[cst.Call, list[cst.ImportFrom]]:
        """Get(<OutputType>, {input1: <Input1Type>, ..inputN: <InputNType>}) =>
        the_rule_to_call(**implicitly(input))"""

        output_type = cst.ensure_type(get.args[0].value, cst.Name).value
        input_dict = cst.ensure_type(get.args[1].value, cst.Dict)
        input_types = {
            element.value.value
            for element in input_dict.elements
            if isinstance(element.value, cst.Name)
        }

        dep = next(
            dep
            for dep in deps
            if dep["output_type"]["name"] == output_type
            and {i["name"] for i in dep["input_types"]} == input_types
        )
        module, new_function = dep["rule_dep"]["module"], dep["rule_dep"]["function"]

        new_call = cst.Call(
            func=cst.Name(new_function),
            args=[
                cst.Arg(
                    value=cst.Call(func=cst.Name("implicitly"), args=[cst.Arg(input_dict)]),
                    star="**",
                )
            ],
        )

        if called_funcdef := cstutil.extract_functiondef_from_module(module, new_function):
            new_call = fix_implicitly_usage(new_call, called_funcdef)

        imports = [cstutil.make_importfrom(module, new_function)]
        return new_call, imports


# ------------------------------------------------------------------------------------------
# Implicity helpers
# ------------------------------------------------------------------------------------------


def fix_implicitly_usage(call: cst.Call, target_func: cst.FunctionDef) -> cst.Call:
    """The CallByNameSyntaxMapper aggressively adds `implicitly` for safety. This function removes
    unnecessary ones, and attempts to cleanup usage.

    Examples:
        find_all_targets(**implicitly()) -> find_all_targets()
        create_pex(**implicitly({req: PexRequest})) -> create_pex(req)
        create_venv_pex(**implicitly({req: PexRequest})) -> create_venv_pex(**implicitly(req))

    Refer to `migrate_call_by_name_test.py` for more examples.

    Parameters:
        call: The replaced `Get` which now uses the migrated call-by-name syntax
        target_func: The target called-by-name function
    """
    call_func_name = cst.ensure_type(call.func, cst.Name).value
    if call_func_name != target_func.name.value:
        return call

    # If there are no `implicitly`s, there is nothing to do
    implicit_calls = m.findall(call, m.Call(func=m.Name("implicitly")))
    if not implicit_calls:
        return call

    # Only handling the 1-arg case (plus implicitly) for now, which is the overwhelming majority of usage
    number_of_call_args = len(call.args)
    if number_of_call_args > 2:
        return call

    # If the target function takes no arguments, then there is nothing to `implicit`
    number_of_target_args = len(target_func.params.params)
    if number_of_target_args == 0:
        return call.with_changes(args=[])

    target_annotations = [
        cst.ensure_type(a, cst.Annotation).annotation
        for a in m.findall(target_func.params, m.Annotation())
    ]
    target_types = [
        target_type for a in target_annotations if (target_type := h.get_full_name_for_node(a))
    ]

    # Positionally compare the target function's annotations with the call's arguments
    # If they match, then there is no need for `implicitly`

    # Check if `implicitly` contains dict - as that needs special handling
    implicit_call = cst.ensure_type(implicit_calls[0], cst.Call)
    if implicit_call.args and isinstance(d := implicit_call.args[0].value, cst.Dict):
        if len(d.elements) > 1:
            # Not handling cases with larger than 1 element
            return call

        element = cst.ensure_type(d.elements[0], cst.DictElement)
        if h.get_full_name_for_node(element.value) == target_types[0]:
            # If arg and target match, we can strip `implicitly` call
            return call.with_changes(args=[cst.Arg(element.key)])
        else:
            # If arg and target don't match, keep `implicitly` call, but remove dict for normal call
            new_arg = cst.Arg(
                cst.Call(cst.Name("implicitly"), args=[cst.Arg(element.key)]), star="**"
            )
            return call.with_changes(args=[new_arg])

    # If the target function takes in the same number of arguments as we've already passed in,
    # and they are of the same type, then we don't need to pass in `implicitly`
    if number_of_call_args - 1 == len(target_types):
        arg = cst.ensure_type(call.args[0].value, cst.Call)
        new_arg = call.args[0].with_changes(comma=cst.MaybeSentinel.DEFAULT)
        if h.get_full_name_for_node(arg.func) == target_types[0]:
            # If arg and target match, we can strip `implicitly` call
            return call.with_changes(args=[new_arg])
        else:
            # If arg and target don't match, keep `implicitly` call
            new_arg = cst.Arg(cst.Call(cst.Name("implicitly"), args=[new_arg]), star="**")
            return call.with_changes(args=[new_arg])

    # If the target function takes in more arguments than we've passed in, then we need to pass in `implicitly`
    # This checks if it should be a trailing implicitly, or if we should wrap the first arg
    if number_of_call_args - 1 < len(target_types):
        arg = cst.ensure_type(call.args[0].value, cst.Call)
        if h.get_full_name_for_node(arg.func) == target_types[0]:
            return call
        else:
            new_arg = call.args[0].with_changes(comma=cst.MaybeSentinel.DEFAULT)
            new_arg = cst.Arg(cst.Call(cst.Name("implicitly"), args=[new_arg]), star="**")
            return call.with_changes(args=[new_arg])

    return call

pantsbuild / pants / 19381742489

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous