Skip to content

Conversation

@frvnkliu
Copy link
Member

@frvnkliu frvnkliu commented Dec 2, 2025

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.

@meta-cla meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Dec 2, 2025
@frvnkliu frvnkliu changed the title Create CodemodCommand Remove/Add Import helpers functions Create CodemodCommand Remove/Add Import helper functions Dec 2, 2025
Comment on lines 72 to 73
obj: Optional[str] = None,
asname: Optional[str] = None,
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can we use the updated str | None syntax instead

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Just updated the commits!

@frvnkliu frvnkliu force-pushed the codemod-command-import-helpers branch from bd06426 to 5a48ffd Compare December 16, 2025 23:05
…r's remove unused import functions in CodemodCommand
… add needed import functions in CodemodCommand
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.
@frvnkliu frvnkliu force-pushed the codemod-command-import-helpers branch from 5a48ffd to 2824627 Compare December 16, 2025 23:18
@frvnkliu frvnkliu merged commit b75343e into Instagram:main Dec 17, 2025
81 of 83 checks passed
@frvnkliu frvnkliu deleted the codemod-command-import-helpers branch December 17, 2025 17:28
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants