Skip to content

Commit

Permalink
refactor: allow scheduled_removals to accept a tuple (#1253)
Browse files Browse the repository at this point in the history
This fixes a TODO
  • Loading branch information
zsol authored Dec 2, 2024
1 parent a36432c commit c05ac74
Showing 1 changed file with 22 additions and 12 deletions.
34 changes: 22 additions & 12 deletions libcst/codemod/commands/rename.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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))

Expand Down Expand Up @@ -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:
Expand Down

0 comments on commit c05ac74

Please sign in to comment.