Skip to content

Commit d77fc49

Browse files
authored
Add docstrings to CodemodCommand import helper methods (#1437)
Add comprehensive docstrings to the remove_unused_import, remove_unused_import_by_node, and add_needed_import helper methods on CodemodCommand class. These docstrings include descriptions, usage examples, and important notes about proper usage.
1 parent ab3e0e3 commit d77fc49

1 file changed

Lines changed: 57 additions & 0 deletions

File tree

libcst/codemod/_command.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,46 @@ def remove_unused_import(
7474
obj: str | None = None,
7575
asname: str | None = None,
7676
) -> None:
77+
"""
78+
Schedule an import to be removed after the codemod completes.
79+
80+
This is a convenience wrapper around the :class:`~libcst.codemod.visitors.RemoveImportsVisitor` static function
81+
:meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import`
82+
that automatically passes the codemod context. The import will only be
83+
removed if it is not referenced elsewhere in the module.
84+
85+
For example, to remove ``from typing import Optional``::
86+
87+
self.remove_unused_import("typing", "Optional")
88+
89+
To remove ``import os``::
90+
91+
self.remove_unused_import("os")
92+
"""
7793
RemoveImportsVisitor.remove_unused_import(self.context, module, obj, asname)
7894

7995
def remove_unused_import_by_node(self, node: CSTNode) -> None:
96+
"""
97+
Schedule removal of all imports referenced by a node and its children.
98+
99+
This is a convenience wrapper around the :class:`~libcst.codemod.visitors.RemoveImportsVisitor` static function
100+
:meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import_by_node`
101+
that automatically passes the codemod context. This is especially useful
102+
when you are removing a node using :func:`~libcst.RemoveFromParent` and want
103+
to clean up any imports that were only used by that node.
104+
105+
For example::
106+
107+
def leave_AnnAssign(
108+
self, original_node: cst.AnnAssign, updated_node: cst.AnnAssign,
109+
) -> cst.RemovalSentinel:
110+
# Remove annotated assignment and clean up imports
111+
self.remove_unused_import_by_node(original_node)
112+
return cst.RemoveFromParent()
113+
114+
Note that you should pass the ``original_node`` rather than ``updated_node``
115+
since scope analysis is computed on the original tree.
116+
"""
80117
RemoveImportsVisitor.remove_unused_import_by_node(self.context, node)
81118

82119
# Lightweight wrappers for AddImportsVisitor static functions
@@ -87,6 +124,26 @@ def add_needed_import(
87124
asname: str | None = None,
88125
relative: int = 0,
89126
) -> None:
127+
"""
128+
Schedule an import to be added after the codemod completes.
129+
130+
This is a convenience wrapper around the :class:`~libcst.codemod.visitors.AddImportsVisitor` static function
131+
:meth:`~libcst.codemod.visitors.AddImportsVisitor.add_needed_import`
132+
that automatically passes the codemod context. The import will only be
133+
added if it does not already exist in the module.
134+
135+
For example, to add ``from typing import Optional``::
136+
137+
self.add_needed_import("typing", "Optional")
138+
139+
To add ``import os``::
140+
141+
self.add_needed_import("os")
142+
143+
To add ``from typing import List as L``::
144+
145+
self.add_needed_import("typing", "List", asname="L")
146+
"""
90147
AddImportsVisitor.add_needed_import(self.context, module, obj, asname, relative)
91148

92149
def transform_module(self, tree: Module) -> Module:

0 commit comments

Comments
 (0)