From 2824627366c4ef4c50cf56138c8b704bada0f570 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Tue, 2 Dec 2025 14:21:44 -0500 Subject: [PATCH 1/3] Create helper functions to abstract away usage of RemoveImportsVisitor's remove unused import functions in CodemodCommand --- libcst/codemod/_command.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/libcst/codemod/_command.py b/libcst/codemod/_command.py index 3671a3fb..78f00dde 100644 --- a/libcst/codemod/_command.py +++ b/libcst/codemod/_command.py @@ -3,12 +3,14 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. # +from __future__ import annotations + import argparse import inspect from abc import ABC, abstractmethod from typing import Dict, Generator, List, Tuple, Type, TypeVar -from libcst import Module +from libcst import CSTNode, Module from libcst.codemod._codemod import Codemod from libcst.codemod._context import CodemodContext from libcst.codemod._visitor import ContextAwareTransformer @@ -65,6 +67,18 @@ def transform_module_impl(self, tree: Module) -> Module: """ ... + # Lightweight wrappers for RemoveImportsVisitor static functions + def remove_unused_import( + self, + module: str, + obj: str | None = None, + asname: str | None = None, + ) -> None: + RemoveImportsVisitor.remove_unused_import(self.context, module, obj, asname) + + def remove_unused_import_by_node(self, node: CSTNode) -> None: + RemoveImportsVisitor.remove_unused_import_by_node(self.context, node) + def transform_module(self, tree: Module) -> Module: # Overrides (but then calls) Codemod's transform_module to provide # a spot where additional supported transforms can be attached and run. From b4848bd74bd14086feab9153cfd46e86f7bddb8d Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Tue, 16 Dec 2025 15:17:52 -0800 Subject: [PATCH 2/3] Create helper functions to abstract away usage of AddImportsVisitor's add needed import functions in CodemodCommand --- libcst/codemod/_command.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/libcst/codemod/_command.py b/libcst/codemod/_command.py index 78f00dde..b7784d30 100644 --- a/libcst/codemod/_command.py +++ b/libcst/codemod/_command.py @@ -79,6 +79,16 @@ def remove_unused_import( def remove_unused_import_by_node(self, node: CSTNode) -> None: RemoveImportsVisitor.remove_unused_import_by_node(self.context, node) + # Lightweight wrappers for AddImportsVisitor static functions + def add_needed_import( + self, + module: str, + obj: str | None = None, + asname: str | None = None, + relative: int = 0, + ) -> None: + AddImportsVisitor.add_needed_import(self.context, module, obj, asname, relative) + def transform_module(self, tree: Module) -> Module: # Overrides (but then calls) Codemod's transform_module to provide # a spot where additional supported transforms can be attached and run. From 463801839b7234ae444adc19d00e21d93f9b05b2 Mon Sep 17 00:00:00 2001 From: Frank Liu Date: Tue, 16 Dec 2025 15:17:52 -0800 Subject: [PATCH 3/3] Add tests for CodemodCommand helper functions Add comprehensive tests for the new helper methods: - remove_unused_import - remove_unused_import_by_node - add_needed_import Tests cover simple cases, from imports, aliased imports, relative imports, and combined add/remove operations. --- libcst/codemod/tests/test_command_helpers.py | 325 +++++++++++++++++++ 1 file changed, 325 insertions(+) create mode 100644 libcst/codemod/tests/test_command_helpers.py diff --git a/libcst/codemod/tests/test_command_helpers.py b/libcst/codemod/tests/test_command_helpers.py new file mode 100644 index 00000000..15c461b6 --- /dev/null +++ b/libcst/codemod/tests/test_command_helpers.py @@ -0,0 +1,325 @@ +# 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. +# +from typing import Union + +import libcst as cst +from libcst.codemod import CodemodTest, VisitorBasedCodemodCommand + + +class TestRemoveUnusedImportHelper(CodemodTest): + """Tests for the remove_unused_import helper method in CodemodCommand.""" + + def test_remove_unused_import_simple(self) -> None: + """ + Test that remove_unused_import helper method works correctly. + """ + + class RemoveBarImport(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Use the helper method to schedule removal + self.remove_unused_import("bar") + + before = """ + import bar + import baz + + def foo() -> None: + pass + """ + after = """ + import baz + + def foo() -> None: + pass + """ + + self.TRANSFORM = RemoveBarImport + self.assertCodemod(before, after) + + def test_remove_unused_import_from_simple(self) -> None: + """ + Test that remove_unused_import helper method works correctly with from imports. + """ + + class RemoveBarFromImport(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Use the helper method to schedule removal + self.remove_unused_import("a.b.c", "bar") + + before = """ + from a.b.c import bar, baz + + def foo() -> None: + baz() + """ + after = """ + from a.b.c import baz + + def foo() -> None: + baz() + """ + + self.TRANSFORM = RemoveBarFromImport + self.assertCodemod(before, after) + + def test_remove_unused_import_with_alias(self) -> None: + """ + Test that remove_unused_import helper method works correctly with aliased imports. + """ + + class RemoveBarAsQuxImport(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Use the helper method to schedule removal + self.remove_unused_import("a.b.c", "bar", "qux") + + before = """ + from a.b.c import bar as qux, baz + + def foo() -> None: + baz() + """ + after = """ + from a.b.c import baz + + def foo() -> None: + baz() + """ + + self.TRANSFORM = RemoveBarAsQuxImport + self.assertCodemod(before, after) + + +class TestRemoveUnusedImportByNodeHelper(CodemodTest): + """Tests for the remove_unused_import_by_node helper method in CodemodCommand.""" + + def test_remove_unused_import_by_node_simple(self) -> None: + """ + Test that remove_unused_import_by_node helper method works correctly. + """ + + class RemoveBarCallAndImport(VisitorBasedCodemodCommand): + METADATA_DEPENDENCIES = ( + cst.metadata.QualifiedNameProvider, + cst.metadata.ScopeProvider, + ) + + def leave_SimpleStatementLine( + self, + original_node: cst.SimpleStatementLine, + updated_node: cst.SimpleStatementLine, + ) -> Union[cst.RemovalSentinel, cst.SimpleStatementLine]: + # Remove any statement that calls bar() + if cst.matchers.matches( + updated_node, + cst.matchers.SimpleStatementLine( + body=[cst.matchers.Expr(cst.matchers.Call())] + ), + ): + call = cst.ensure_type(updated_node.body[0], cst.Expr).value + if cst.matchers.matches( + call, cst.matchers.Call(func=cst.matchers.Name("bar")) + ): + # Use the helper method to remove imports referenced by this node + self.remove_unused_import_by_node(original_node) + return cst.RemoveFromParent() + return updated_node + + before = """ + from foo import bar, baz + + def fun() -> None: + bar() + baz() + """ + after = """ + from foo import baz + + def fun() -> None: + baz() + """ + + self.TRANSFORM = RemoveBarCallAndImport + self.assertCodemod(before, after) + + +class TestAddNeededImportHelper(CodemodTest): + """Tests for the add_needed_import helper method in CodemodCommand.""" + + def test_add_needed_import_simple(self) -> None: + """ + Test that add_needed_import helper method works correctly. + """ + + class AddBarImport(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Use the helper method to schedule import addition + self.add_needed_import("bar") + + before = """ + def foo() -> None: + pass + """ + after = """ + import bar + + def foo() -> None: + pass + """ + + self.TRANSFORM = AddBarImport + self.assertCodemod(before, after) + + def test_add_needed_import_from_simple(self) -> None: + """ + Test that add_needed_import helper method works correctly with from imports. + """ + + class AddBarFromImport(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Use the helper method to schedule import addition + self.add_needed_import("a.b.c", "bar") + + before = """ + def foo() -> None: + pass + """ + after = """ + from a.b.c import bar + + def foo() -> None: + pass + """ + + self.TRANSFORM = AddBarFromImport + self.assertCodemod(before, after) + + def test_add_needed_import_with_alias(self) -> None: + """ + Test that add_needed_import helper method works correctly with aliased imports. + """ + + class AddBarAsQuxImport(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Use the helper method to schedule import addition + self.add_needed_import("a.b.c", "bar", "qux") + + before = """ + def foo() -> None: + pass + """ + after = """ + from a.b.c import bar as qux + + def foo() -> None: + pass + """ + + self.TRANSFORM = AddBarAsQuxImport + self.assertCodemod(before, after) + + def test_add_needed_import_relative(self) -> None: + """ + Test that add_needed_import helper method works correctly with relative imports. + """ + + class AddRelativeImport(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Use the helper method to schedule relative import addition + self.add_needed_import("c", "bar", relative=2) + + before = """ + def foo() -> None: + pass + """ + after = """ + from ..c import bar + + def foo() -> None: + pass + """ + + self.TRANSFORM = AddRelativeImport + self.assertCodemod(before, after) + + +class TestCombinedHelpers(CodemodTest): + """Tests for combining add_needed_import and remove_unused_import helper methods.""" + + def test_add_and_remove_imports(self) -> None: + """ + Test that both helper methods work correctly when used together. + """ + + class ReplaceBarWithBaz(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Add new import and remove old one + self.add_needed_import("new_module", "baz") + self.remove_unused_import("old_module", "bar") + + before = """ + from other_module import qux + from old_module import bar + + def foo() -> None: + pass + """ + after = """ + from other_module import qux + from new_module import baz + + def foo() -> None: + pass + """ + + self.TRANSFORM = ReplaceBarWithBaz + self.assertCodemod(before, after) + + def test_add_and_remove_same_import(self) -> None: + """ + Test that both helper methods work correctly when used together. + """ + + class AddAndRemoveBar(VisitorBasedCodemodCommand): + def visit_Module(self, node: cst.Module) -> None: + # Add new import and remove old one + self.add_needed_import("hello_module", "bar") + self.remove_unused_import("hello_module", "bar") + + self.TRANSFORM = AddAndRemoveBar + + before = """ + from other_module import baz + + def foo() -> None: + pass + """ + # Should remain unchanged + self.assertCodemod(before, before) + + before = """ + from other_module import baz + from hello_module import bar + + def foo() -> None: + bar.func() + """ + self.assertCodemod(before, before) + + before = """ + from other_module import baz + from hello_module import bar + + def foo() -> None: + pass + """ + + after = """ + from other_module import baz + + def foo() -> None: + pass + """ + self.assertCodemod(before, after)