LibCST icon indicating copy to clipboard operation
LibCST copied to clipboard

Create CodemodCommand Remove/Add Import helper functions

Open frvnkliu opened this issue 3 months ago • 0 comments

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:

  1. RemoveImportsVisitor.remove_unused_import
  2. RemoveImportsVisitor.remove_unused_import_by_node
  3. 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:

  1. remove_unused_import
  2. remove_unused_import_by_node
  3. 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

  1. Format your code

    uv run poe format

  2. Run the type checker

    uv run poe typecheck

  3. Test your changes

    uv run poe test

  4. 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.

frvnkliu avatar Dec 02 '25 20:12 frvnkliu