Create CodemodCommand Remove/Add Import helper functions
Summary
One of the key features of the CodemodCommand class is that it abstracts away the final run of AddImportsVisitor and RemoveImportsVisitor in transform_module (Code Source). However, currently developers still need to utilize the static functions:
-
RemoveImportsVisitor.remove_unused_import -
RemoveImportsVisitor.remove_unused_import_by_node -
AddImportsVisitor.add_needed_import
To actually manage the import states in the self.context: CodemodContext. I propose defining these three new instance functions on CodemodCommand:
-
remove_unused_import -
remove_unused_import_by_node -
add_needed_import
These new lightweight helper wrapper functions defined on CodemodCommand will help to hopefully completely abstract away RemoveImportsVisitor and AddImportsVisitor.
[Example] Remove + Add Imports
Before:
from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor
class RemoveAddImportExampleCommand(VisitorBasedCodemodCommand):
def visit_Module(self, node: cst.Module) -> None:
RemoveImportsVisitor.remove_unused_import(self.context, "a.b.c", "bar")
AddImportsVisitor.add_needed_import(self.context, "hello_module", "bar")
After:
# No more extra imports :)
class RemoveAddImportExampleCommand(VisitorBasedCodemodCommand):
def visit_Module(self, node: cst.Module) -> None:
self.remove_unused_import("a.b.c", "bar") # no need to pass in self.context anymore :)
self.add_needed_import("hello_module", "bar")
Test Plan
- Format your code
uv run poe format
- Run the type checker
uv run poe typecheck
- Test your changes
uv run poe test
- Check linters
uv run poe lint
Specific Unit Tests found in:
libcst/codemod/tests/test_command_helpers.py
and can be run using
uv run poe test -k test_command_helpers
[IG][Meta] Would be nice to have this to lower barrier of entry for using CodemodCommand.