From f7bb0cec9392b2883d89e0522d26de820985b978 Mon Sep 17 00:00:00 2001 From: Zsolt Dollenstein Date: Sat, 30 Nov 2024 10:55:27 +0000 Subject: [PATCH] refactor: allow scheduled_removals to accept a tuple This fixes a TODO --- libcst/codemod/commands/rename.py | 34 ++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 12 deletions(-) diff --git a/libcst/codemod/commands/rename.py b/libcst/codemod/commands/rename.py index aad4cea6..8b6302c9 100644 --- a/libcst/codemod/commands/rename.py +++ b/libcst/codemod/commands/rename.py @@ -103,15 +103,20 @@ def as_name(self, value: Optional[Tuple[str, str]]) -> None: self.context.scratch["as_name"] = value @property - def scheduled_removals(self) -> Set[cst.CSTNode]: + def scheduled_removals( + self, + ) -> Set[Union[cst.CSTNode, Tuple[str, Optional[str], Optional[str]]]]: """A set of nodes that have been renamed to help with the cleanup of now potentially unused - imports, during import cleanup in `leave_Module`.""" + imports, during import cleanup in `leave_Module`. Can also contain tuples that can be passed + directly to RemoveImportsVisitor.remove_unused_import().""" if "scheduled_removals" not in self.context.scratch: self.context.scratch["scheduled_removals"] = set() return self.context.scratch["scheduled_removals"] @scheduled_removals.setter - def scheduled_removals(self, value: Set[cst.CSTNode]) -> None: + def scheduled_removals( + self, value: Set[Union[cst.CSTNode, Tuple[str, Optional[str], Optional[str]]]] + ) -> None: self.context.scratch["scheduled_removals"] = value @property @@ -169,11 +174,13 @@ def leave_Import( and import_alias.asname is not None ): self.bypass_import = True - # TODO: put this into self.scheduled_removals - RemoveImportsVisitor.remove_unused_import( - self.context, - import_alias.evaluated_name, - asname=import_alias.evaluated_alias, + # Add removal tuple instead of calling directly + self.scheduled_removals.add( + ( + import_alias.evaluated_name, + None, + import_alias.evaluated_alias, + ) ) new_names.append(import_alias.with_changes(asname=None)) @@ -307,10 +314,13 @@ def leave_Attribute( 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 - ) + for removal in self.scheduled_removals: + if isinstance(removal, tuple): + RemoveImportsVisitor.remove_unused_import( + self.context, removal[0], removal[1], removal[2] + ) + else: + RemoveImportsVisitor.remove_unused_import_by_node(self.context, removal) # 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: