@@ -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