path: root/scripts/codemods/rename_pyqt.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# pyre-strict

"""libCST codemod to rename imports and references to them."""

import argparse
from fnmatch import fnmatch
from typing import Callable, Optional, Sequence, Set, Tuple, Union

import libcst as cst
from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor
from libcst.helpers import get_full_name_for_node
from libcst.metadata import QualifiedNameProvider


def leave_import_decorator(
    method: Callable[..., Union[cst.Import, cst.ImportFrom]]
) -> Callable[..., Union[cst.Import, cst.ImportFrom]]:
    # We want to record any 'as name' that is relevant but only after we leave the corresponding Import/ImportFrom node since
    # we don't want the 'as name' to interfere with children 'Name' and 'Attribute' nodes.
    def wrapper(
        self: "RenameCommand",
        original_node: Union[cst.Import, cst.ImportFrom],
        updated_node: Union[cst.Import, cst.ImportFrom],
    ) -> Union[cst.Import, cst.ImportFrom]:
        updated_node = method(self, original_node, updated_node)
        if original_node != updated_node:
            self.record_asname(original_node)
        return updated_node

    return wrapper


class RenameCommand(VisitorBasedCodemodCommand):
    """Rename all instances of a local or imported object."""

    DESCRIPTION: str = "Rename all instances of a local or imported object."

    METADATA_DEPENDENCIES = (QualifiedNameProvider,)

    @staticmethod
    def add_args(arg_parser: argparse.ArgumentParser) -> None:
        arg_parser.add_argument(
            "--old_name",
            dest="old_name",
            required=True,
            help="Full dotted name of object to rename. Eg: `foo.bar.baz`",
        )

        arg_parser.add_argument(
            "--new_name",
            dest="new_name",
            required=True,
            help=(
                "Full dotted name of replacement object. You may provide a single-colon-delimited name to specify how you want the new import to be structured."
                + "\nEg: `foo:bar.baz` will be translated to `from foo import bar`."
                + "\nIf no ':' character is provided, the import statement will default to `from foo.bar import baz` for a `new_name` value of `foo.bar.baz`"
                + " or simply replace the old import on the spot if the old import is an exact match."
            ),
        )

    def __init__(self, context: CodemodContext, old_name: str, new_name: str) -> None:
        super().__init__(context)

        new_module, has_colon, new_mod_or_obj = new_name.rpartition(":")
        # Exit early if improperly formatted args.
        if ":" in new_module:
            raise ValueError("Error: `new_name` should contain at most one colon.")
        if ":" in old_name:
            raise ValueError("Error: `old_name` should not contain any colons.")

        if not has_colon or not new_module:
            new_module, _, new_mod_or_obj = new_name.rpartition(".")

        self.new_name: str = new_name.replace(":", ".").strip(".")
        self.new_module: str = new_module.replace(":", ".").strip(".")
        self.new_mod_or_obj: str = new_mod_or_obj if new_mod_or_obj != "*" else None

        # If `new_name` contains a single colon at the end, then we assume the user wants the import
        # to be structured as 'import new_name'. So both self.new_mod_or_obj and self.old_mod_or_obj
        # will be empty in this case.
        if not self.new_mod_or_obj:
            old_module = old_name
            old_mod_or_obj = ""
        else:
            old_module, _, old_mod_or_obj = old_name.rpartition(".")

        self.old_name: str = old_name
        self.old_module: str = old_module
        self.old_mod_or_obj: str = old_mod_or_obj

        self.as_name: Optional[Tuple[str, str]] = None

        # A set of nodes that have been renamed to help with the cleanup of now potentially unused
        # imports, during import cleanup in `leave_Module`.
        self.scheduled_removals: Set[cst.CSTNode] = set()
        # If an import has been renamed while inside an `Import` or `ImportFrom` node, we want to flag
        # this so that we do not end up with two of the same import.
        self.bypass_import = False

    def visit_Import(self, node: cst.Import) -> None:
        for import_alias in node.names:
            alias_name = get_full_name_for_node(import_alias.name)
            if alias_name is not None:
                if alias_name == self.old_name or alias_name.startswith(
                    self.old_name + "."
                ):
                    # If the import statement is exactly equivalent to the old name, or we are renaming a top-level module of the import,
                    # it will be taken care of in `leave_Name` or `leave_Attribute` when visiting the Name and Attribute children of this Import.
                    self.bypass_import = True

    @leave_import_decorator
    def leave_Import(
        self, original_node: cst.Import, updated_node: cst.Import
    ) -> cst.Import:
        new_names = []
        for import_alias in updated_node.names:
            import_alias_name = import_alias.name
            import_alias_full_name = get_full_name_for_node(import_alias_name)
            if import_alias_full_name is None:
                raise Exception("Could not parse full name for ImportAlias.name node.")

            if isinstance(import_alias_name, cst.Name) and self.old_name.startswith(
                import_alias_full_name + "."
            ):
                # Might, be in use elsewhere in the code, so schedule a potential removal, and add another alias.
                new_names.append(import_alias)
                self.scheduled_removals.add(original_node)
                new_names.append(
                    cst.ImportAlias(
                        name=cst.Name(
                            value=self.gen_replacement_module(import_alias_full_name)
                        )
                    )
                )
                self.bypass_import = True
            elif isinstance(
                import_alias_name, cst.Attribute
            ) and self.old_name.startswith(import_alias_full_name + "."):
                # Same idea as above.
                new_names.append(import_alias)
                self.scheduled_removals.add(original_node)
                new_name_node: Union[
                    cst.Attribute, cst.Name
                ] = self.gen_name_or_attr_node(
                    self.gen_replacement_module(import_alias_full_name)
                )
                new_names.append(cst.ImportAlias(name=new_name_node))
                self.bypass_import = True
            else:
                new_names.append(import_alias)

        return updated_node.with_changes(names=new_names)

    def visit_ImportFrom(self, node: cst.ImportFrom) -> None:
        module = node.module
        if module is None:
            return
        imported_module_name = get_full_name_for_node(module)
        if imported_module_name is None:
            return
        if imported_module_name == self.old_name or imported_module_name.startswith(
            self.old_name + "."
        ):
            # If the imported module is exactly equivalent to the old name or we are renaming a parent module of the current module,
            # it will be taken care of in `leave_Name` or `leave_Attribute` when visiting the children of this ImportFrom.
            self.bypass_import = True

    @leave_import_decorator
    def leave_ImportFrom(
        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
    ) -> cst.ImportFrom:
        module = updated_node.module
        if module is None:
            return updated_node
        imported_module_name = get_full_name_for_node(module)
        names = original_node.names

        if imported_module_name is None or not isinstance(names, Sequence):
            return updated_node

        else:
            new_names = []
            for import_alias in names:
                alias_name = get_full_name_for_node(import_alias.name)
                if alias_name is not None:
                    qual_name = f"{imported_module_name}.{alias_name}"
                    if fnmatch(qual_name, self.old_name):

                        replacement_module = self.gen_replacement_module(
                            imported_module_name
                        )
                        replacement_obj = self.gen_replacement(alias_name)
                        if not replacement_obj:
                            # The user has requested an `import` statement rather than an `from ... import`.
                            # This will be taken care of in `leave_Module`, in the meantime, schedule for potential removal.
                            new_names.append(import_alias)
                            self.scheduled_removals.add(original_node)
                            continue

                        new_import_alias_name: Union[
                            cst.Attribute, cst.Name
                        ] = self.gen_name_or_attr_node(replacement_obj)
                        # Rename on the spot only if this is the only imported name under the module.
                        if len(names) == 1:
                            print(f"just one {new_import_alias_name}")
                            self.bypass_import = True
                            return updated_node.with_changes(
                                module=cst.parse_expression(replacement_module),
                                names=(cst.ImportAlias(name=new_import_alias_name),),
                            )
                        # Or if the module name is to stay the same.
                        elif replacement_module == imported_module_name:
                            self.bypass_import = True
                            new_names.append(
                                cst.ImportAlias(name=new_import_alias_name)
                            )
                    else:
                        if self.old_name.startswith(qual_name + "."):
                            # This import might be in use elsewhere in the code, so schedule a potential removal.
                            self.scheduled_removals.add(original_node)
                        new_names.append(import_alias)

            return updated_node.with_changes(names=new_names)
        return updated_node

    def has_name_wildcard(self, original_node, old_name):
        qualified_names = self.get_metadata(QualifiedNameProvider, original_node, set())
        if isinstance(old_name, str):
            return any(fnmatch(qn.name, old_name) for qn in qualified_names)
        else:
            return any(fnmatch(qn == old_name) for qn in qualified_names)

    def leave_Name(
        self, original_node: cst.Name, updated_node: cst.Name
    ) -> Union[cst.Attribute, cst.Name]:
        full_name_for_node: str = original_node.value
        full_replacement_name = self.gen_replacement(full_name_for_node)

        # If a node has no associated QualifiedName, we are still inside an import statement.
        inside_import_statement: bool = not self.get_metadata(
            QualifiedNameProvider, original_node, set()
        )
        #qns = self.get_metadata(
        #    QualifiedNameProvider, original_node, set()
        #)
        #print(original_node)
        #print(f"{full_name_for_node} vs {self.old_name} and {qns}")
        if self.has_name_wildcard(original_node, self.old_name) or (
            inside_import_statement and full_replacement_name == self.new_name
        ) or (
            fnmatch(self.old_name, full_name_for_node)
        ):
            if not full_replacement_name:
                full_replacement_name = self.new_mod_or_obj
            if not inside_import_statement:
                self.scheduled_removals.add(original_node)
            if full_replacement_name.endswith('.*'):
                full_replacement_name, _, _ = full_replacement_name.rpartition('.')
                full_replacement_name = ".".join((full_replacement_name, full_name_for_node))
            #print(f'yes a {full_replacement_name} from {self.new_mod_or_obj}')
            return self.gen_name_or_attr_node(full_replacement_name)
        else:
            #print('no')
            pass

        return updated_node

    def leave_Attribute(
        self, original_node: cst.Attribute, updated_node: cst.Attribute
    ) -> Union[cst.Name, cst.Attribute]:
        full_name_for_node = get_full_name_for_node(original_node)
        if full_name_for_node is None:
            raise Exception("Could not parse full name for Attribute node.")
        full_replacement_name = self.gen_replacement(full_name_for_node)

        # If a node has no associated QualifiedName, we are still inside an import statement.
        inside_import_statement: bool = not self.get_metadata(
            QualifiedNameProvider, original_node, set()
        )
        if self.has_name_wildcard(
            original_node,
            self.old_name,
        ) or (inside_import_statement and full_replacement_name == self.new_name):
            new_value, new_attr = self.new_module, self.new_mod_or_obj
            if not inside_import_statement:
                self.scheduled_removals.add(original_node.value)
            print(f"{full_replacement_name} vs {self.new_name}")
            print(f"{new_value} and {new_attr}")
            print(f"old {self.old_name} {original_node} {original_node.value}")
            qualified_names = self.get_metadata(QualifiedNameProvider, original_node, set())
            print(f"qns {qualified_names}")
            if full_replacement_name == self.new_name:
                return updated_node.with_changes(
                    value=cst.parse_expression(new_value),
                    attr=cst.Name(value=new_attr.rstrip(".")),
                )

            if new_attr.endswith('.*'):
                new_attr, _, _ = new_attr.rpartition('.')
                new_attr = ".".join((new_attr, full_name_for_node))
            return self.gen_name_or_attr_node(new_attr)

        return updated_node

    def leave_Module(
        self, original_node: cst.Module, updated_node: cst.Module
    ) -> cst.Module:
        for removal_node in self.scheduled_removals:
            RemoveImportsVisitor.remove_unused_import_by_node(
                self.context, removal_node
            )
        # If bypass_import is False, we know that no import statements were directly renamed, and the fact
        # that we have any `self.scheduled_removals` tells us we encountered a matching `old_name` in the code.
        if not self.bypass_import and self.scheduled_removals:
            if self.new_module:
                new_obj: Optional[str] = (
                    self.new_mod_or_obj.split(".")[0] if self.new_mod_or_obj else None
                )
                AddImportsVisitor.add_needed_import(
                    self.context, module=self.new_module, obj=new_obj
                )
        return updated_node

    def gen_replacement(self, original_name: str) -> str:
        module_as_name = self.as_name
        if module_as_name is not None:
            if original_name == module_as_name[0]:
                original_name = module_as_name[1]
            elif original_name.startswith(module_as_name[0] + "."):
                original_name = original_name.replace(
                    module_as_name[0] + ".", module_as_name[1] + ".", 1
                )

        if False and self.old_mod_or_obj == "*":
            print(f"wildcard to {original_name}")
            #return original_name  # QIcon
            #return self.new_mod_or_obj  # QtGui.*
            return self.new_name
        elif original_name == self.old_mod_or_obj:
            #print(f"orig = old {self.new_mod_or_obj}")
            return self.new_mod_or_obj
        elif original_name == ".".join([self.old_module, self.old_mod_or_obj]):
            #print(f"orig = full.old {self.old_module}.{self.new_mod_or_obj}")
            return self.new_name
        elif original_name.endswith("." + self.old_mod_or_obj):
            #print(f"orig.endswith(old) {self.new_mod_or_obj}")
            return self.new_mod_or_obj
        else:
            new = self.gen_replacement_module(original_name)
            #print(f"{original_name} -> {new}")
            return new

    def gen_replacement_module(self, original_module: str) -> str:
        return self.new_module if fnmatch(original_module, self.old_module) else ""

    def gen_name_or_attr_node(
        self, dotted_expression: str
    ) -> Union[cst.Attribute, cst.Name]:
        name_or_attr_node: cst.BaseExpression = cst.parse_expression(dotted_expression)
        if not isinstance(name_or_attr_node, (cst.Name, cst.Attribute)):
            raise Exception(
                "`parse_expression()` on dotted path returned non-Attribute-or-Name."
            )
        return name_or_attr_node

    def record_asname(self, original_node: Union[cst.Import, cst.ImportFrom]) -> None:
        # Record the import's `as` name if it has one, and set the attribute mapping.
        names = original_node.names
        if not isinstance(names, Sequence):
            return
        for import_alias in names:
            alias_name = get_full_name_for_node(import_alias.name)
            if isinstance(original_node, cst.ImportFrom):
                module = original_node.module
                if module is None:
                    return
                module_name = get_full_name_for_node(module)
                if module_name is None:
                    return
                qual_name = f"{module_name}.{alias_name}"
            else:
                qual_name = alias_name
            if qual_name is not None and alias_name is not None:
                if fnmatch(qual_name, self.old_name) or self.old_name.startswith(
                    qual_name + "."
                ):
                    as_name_optional = import_alias.asname
                    as_name_node = (
                        as_name_optional.name if as_name_optional is not None else None
                    )
                    #print(f"{qual_name} as {as_name_node}")
                    if as_name_node is not None and isinstance(
                        as_name_node, (cst.Name, cst.Attribute)
                    ):
                        full_as_name = get_full_name_for_node(as_name_node)
                        if full_as_name is not None:
                            self.as_name = (full_as_name, alias_name)