Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2#
3# This source code is licensed under the MIT license found in the
4# LICENSE file in the root directory of this source tree.
5#
6from typing import Any, Dict, Iterable, List, Optional, Sequence, Set, Tuple, Union
8import libcst as cst
9from libcst import CSTLogicError
10from libcst.codemod._context import CodemodContext
11from libcst.codemod._visitor import ContextAwareTransformer, ContextAwareVisitor
12from libcst.codemod.visitors._gather_unused_imports import GatherUnusedImportsVisitor
13from libcst.helpers import (
14 get_absolute_module_from_package_for_import,
15 get_full_name_for_node,
16)
17from libcst.metadata import Assignment, ProviderT, ScopeProvider
20class RemovedNodeVisitor(ContextAwareVisitor):
21 def _remove_imports_from_import_stmt(
22 self, local_name: str, import_node: cst.Import
23 ) -> None:
24 for import_alias in import_node.names:
25 if import_alias.evaluated_alias is None:
26 prefix = import_alias.evaluated_name
27 else:
28 prefix = import_alias.evaluated_alias
30 if local_name == prefix or local_name.startswith(f"{prefix}."):
31 RemoveImportsVisitor.remove_unused_import(
32 self.context,
33 import_alias.evaluated_name,
34 asname=import_alias.evaluated_alias,
35 )
37 def _remove_imports_from_importfrom_stmt(
38 self, local_name: str, import_node: cst.ImportFrom
39 ) -> None:
40 names = import_node.names
41 if isinstance(names, cst.ImportStar):
42 # We don't handle removing this, so ignore it.
43 return
45 module_name = get_absolute_module_from_package_for_import(
46 self.context.full_package_name, import_node
47 )
48 if module_name is None:
49 raise ValueError("Cannot look up absolute module from relative import!")
51 # We know any local names will refer to this as an alias if
52 # there is one, and as the original name if there is not one
53 for import_alias in names:
54 if import_alias.evaluated_alias is None:
55 prefix = import_alias.evaluated_name
56 else:
57 prefix = import_alias.evaluated_alias
59 if local_name == prefix or local_name.startswith(f"{prefix}."):
60 RemoveImportsVisitor.remove_unused_import(
61 self.context,
62 module_name,
63 obj=import_alias.evaluated_name,
64 asname=import_alias.evaluated_alias,
65 )
67 def _visit_name_attr_alike(self, node: Union[cst.Name, cst.Attribute]) -> None:
68 # Look up the local name of this node.
69 local_name = get_full_name_for_node(node)
70 if local_name is None:
71 return
73 # Look up the scope for this node, remove the import that caused it to exist.
74 metadata_wrapper = self.context.wrapper
75 if metadata_wrapper is None:
76 raise ValueError(
77 "Cannot look up import, metadata is not computed for node!"
78 )
79 scope_provider = metadata_wrapper.resolve(ScopeProvider)
80 try:
81 scope = scope_provider[node]
82 if scope is None:
83 # This object has no scope, so we can't remove it.
84 return
85 except KeyError:
86 # This object has no scope, so we can't remove it.
87 return
89 while True:
90 for assignment in scope.assignments[node] or set():
91 # We only care about non-builtins.
92 if isinstance(assignment, Assignment):
93 import_node = assignment.node
94 if isinstance(import_node, cst.Import):
95 self._remove_imports_from_import_stmt(local_name, import_node)
96 elif isinstance(import_node, cst.ImportFrom):
97 self._remove_imports_from_importfrom_stmt(
98 local_name, import_node
99 )
101 if scope is scope.parent:
102 break
103 scope = scope.parent
105 def visit_Name(self, node: cst.Name) -> None:
106 self._visit_name_attr_alike(node)
108 def visit_Attribute(self, node: cst.Attribute) -> None:
109 self._visit_name_attr_alike(node)
112class RemoveImportsVisitor(ContextAwareTransformer):
113 """
114 Attempt to remove given imports from a module, dependent on whether there are
115 any uses of the imported objects. Given a :class:`~libcst.codemod.CodemodContext`
116 and a sequence of tuples specifying a module to remove as a string. Optionally
117 an object being imported from that module and optionally an alias assigned to
118 that imported object, ensures that that import no longer exists as long as there
119 are no remaining references.
121 Note that static analysis is able to determine safely whether an import is still
122 needed given a particular module, but it is currently unable to determine whether
123 an imported object is re-exported and used inside another module unless that
124 object appears in an ``__any__`` list.
126 This is one of the transforms that is available automatically to you when running
127 a codemod. To use it in this manner, import
128 :class:`~libcst.codemod.visitors.RemoveImportsVisitor` and then call the static
129 :meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import` method,
130 giving it the current context (found as ``self.context`` for all subclasses of
131 :class:`~libcst.codemod.Codemod`), the module you wish to remove and
132 optionally an object you wish to stop importing as well as an alias that the
133 object is currently assigned to.
135 For example::
137 RemoveImportsVisitor.remove_unused_import(self.context, "typing", "Optional")
139 This will remove any ``from typing import Optional`` that exists in the module
140 as long as there are no uses of ``Optional`` in that module.
142 As another example::
144 RemoveImportsVisitor.remove_unused_import(self.context, "typing")
146 This will remove any ``import typing`` that exists in the module, as long as
147 there are no references to ``typing`` in that module, including references
148 such as ``typing.Optional``.
150 Additionally, :class:`~libcst.codemod.visitors.RemoveImportsVisitor` includes
151 a convenience function
152 :meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import_by_node`
153 which will attempt to schedule removal of all imports referenced in that node
154 and its children. This is especially useful inside transforms when you are going
155 to remove a node using :func:`~libcst.RemoveFromParent` to get rid of a node.
157 For example::
159 def leave_AnnAssign(
160 self, original_node: cst.AnnAssign, updated_node: cst.AnnAssign,
161 ) -> cst.RemovalSentinel:
162 # Remove all annotated assignment statements, clean up imports.
163 RemoveImportsVisitor.remove_unused_import_by_node(self.context, original_node)
164 return cst.RemovalFromParent()
166 This will remove all annotated assignment statements from a module as well
167 as clean up any imports that were only referenced in those assignments. Note
168 that we pass the ``original_node`` to the helper function as it uses scope analysis
169 under the hood which is only computed on the original tree.
171 Note that this is a subclass of :class:`~libcst.CSTTransformer` so it is
172 possible to instantiate it and pass it to a :class:`~libcst.Module`
173 :meth:`~libcst.CSTNode.visit` method. However, it is far easier to use
174 the automatic transform feature of :class:`~libcst.codemod.CodemodCommand`
175 and schedule an import to be added by calling
176 :meth:`~libcst.codemod.visitors.RemoveImportsVisitor.remove_unused_import`
178 """
180 CONTEXT_KEY = "RemoveImportsVisitor"
181 METADATA_DEPENDENCIES: Tuple[ProviderT] = (
182 *GatherUnusedImportsVisitor.METADATA_DEPENDENCIES,
183 )
185 @staticmethod
186 def _get_imports_from_context(
187 context: CodemodContext,
188 ) -> List[Tuple[str, Optional[str], Optional[str]]]:
189 unused_imports = context.scratch.get(RemoveImportsVisitor.CONTEXT_KEY, [])
190 if not isinstance(unused_imports, list):
191 raise CSTLogicError("Logic error!")
192 return unused_imports
194 @staticmethod
195 def remove_unused_import(
196 context: CodemodContext,
197 module: str,
198 obj: Optional[str] = None,
199 asname: Optional[str] = None,
200 ) -> None:
201 """
202 Schedule an import to be removed in a future invocation of this class by
203 updating the ``context`` to include the ``module`` and optionally ``obj``
204 which is currently imported as well as optionally ``alias`` that the
205 imported ``module`` or ``obj`` is aliased to. When subclassing from
206 :class:`~libcst.codemod.CodemodCommand`, this will be performed for you
207 after your transform finishes executing. If you are subclassing from a
208 :class:`~libcst.codemod.Codemod` instead, you will need to call the
209 :meth:`~libcst.codemod.Codemod.transform_module` method on the module
210 under modification with an instance of this class after performing your
211 transform. Note that if the particular ``module`` or ``obj`` you are
212 requesting to remove is still in use somewhere in the current module
213 at the time of executing :meth:`~libcst.codemod.Codemod.transform_module`
214 on an instance of :class:`~libcst.codemod.visitors.AddImportsVisitor`,
215 this will perform no action in order to avoid removing an in-use import.
216 """
218 unused_imports = RemoveImportsVisitor._get_imports_from_context(context)
219 unused_imports.append((module, obj, asname))
220 context.scratch[RemoveImportsVisitor.CONTEXT_KEY] = unused_imports
222 @staticmethod
223 def remove_unused_import_by_node(
224 context: CodemodContext, node: cst.CSTNode
225 ) -> None:
226 """
227 Schedule any imports referenced by ``node`` or one of its children
228 to be removed in a future invocation of this class by updating the
229 ``context`` to include the ``module``, ``obj`` and ``alias`` for each
230 import in question. When subclassing from
231 :class:`~libcst.codemod.CodemodCommand`, this will be performed for you
232 after your transform finishes executing. If you are subclassing from a
233 :class:`~libcst.codemod.Codemod` instead, you will need to call the
234 :meth:`~libcst.codemod.Codemod.transform_module` method on the module
235 under modification with an instance of this class after performing your
236 transform. Note that all imports that are referenced by this ``node``
237 or its children will only be removed if they are not in use at the time
238 of exeucting :meth:`~libcst.codemod.Codemod.transform_module`
239 on an instance of :class:`~libcst.codemod.visitors.AddImportsVisitor`
240 in order to avoid removing an in-use import.
241 """
243 # Special case both Import and ImportFrom so they can be
244 # directly removed here.
245 if isinstance(node, cst.Import):
246 for import_alias in node.names:
247 RemoveImportsVisitor.remove_unused_import(
248 context,
249 import_alias.evaluated_name,
250 asname=import_alias.evaluated_alias,
251 )
252 elif isinstance(node, cst.ImportFrom):
253 names = node.names
254 if isinstance(names, cst.ImportStar):
255 # We don't handle removing this, so ignore it.
256 return
257 module_name = get_absolute_module_from_package_for_import(
258 context.full_package_name, node
259 )
260 if module_name is None:
261 raise ValueError("Cannot look up absolute module from relative import!")
262 for import_alias in names:
263 RemoveImportsVisitor.remove_unused_import(
264 context,
265 module_name,
266 obj=import_alias.evaluated_name,
267 asname=import_alias.evaluated_alias,
268 )
269 else:
270 # Look up all children that could have been imported. Any that
271 # we find will be scheduled for removal.
272 node.visit(RemovedNodeVisitor(context))
274 def __init__(
275 self,
276 context: CodemodContext,
277 unused_imports: Sequence[Tuple[str, Optional[str], Optional[str]]] = (),
278 ) -> None:
279 # Allow for instantiation from either a context (used when multiple transforms
280 # get chained) or from a direct instantiation.
281 super().__init__(context)
283 all_unused_imports: List[Tuple[str, Optional[str], Optional[str]]] = [
284 *RemoveImportsVisitor._get_imports_from_context(context),
285 *unused_imports,
286 ]
287 self.unused_module_imports: Dict[str, Optional[str]] = {
288 module: alias for module, obj, alias in all_unused_imports if obj is None
289 }
290 self.unused_obj_imports: Dict[str, Set[Tuple[str, Optional[str]]]] = {}
291 for module, obj, alias in all_unused_imports:
292 if obj is None:
293 continue
294 if module not in self.unused_obj_imports:
295 self.unused_obj_imports[module] = set()
296 self.unused_obj_imports[module].add((obj, alias))
297 self._unused_imports: Dict[
298 cst.ImportAlias, Union[cst.Import, cst.ImportFrom]
299 ] = {}
301 def visit_Module(self, node: cst.Module) -> None:
302 visitor = GatherUnusedImportsVisitor(self.context)
303 node.visit(visitor)
304 self._unused_imports = {k: v for (k, v) in visitor.unused_imports}
306 def leave_Import(
307 self, original_node: cst.Import, updated_node: cst.Import
308 ) -> Union[cst.Import, cst.RemovalSentinel]:
309 names_to_keep = []
310 for import_alias in original_node.names:
311 if import_alias.evaluated_name not in self.unused_module_imports:
312 # This is a keeper since we aren't removing it
313 names_to_keep.append(import_alias)
314 continue
316 if (
317 import_alias.evaluated_alias
318 != self.unused_module_imports[import_alias.evaluated_name]
319 ):
320 # This is a keeper since the alias does not match
321 # what we are looking for.
322 names_to_keep.append(import_alias)
323 continue
325 # Now that we know we want to remove this module, figure out if
326 # there are any live references to it.
327 if import_alias not in self._unused_imports:
328 names_to_keep.append(import_alias)
329 continue
331 # no changes
332 if names_to_keep == original_node.names:
333 return updated_node
335 # Now, either remove this statement or remove the imports we are
336 # deleting from this statement.
337 if len(names_to_keep) == 0:
338 return cst.RemoveFromParent()
340 if names_to_keep[-1] != original_node.names[-1]:
341 # Remove trailing comma in order to not mess up import statements.
342 names_to_keep = [
343 *names_to_keep[:-1],
344 names_to_keep[-1].with_changes(comma=cst.MaybeSentinel.DEFAULT),
345 ]
346 return updated_node.with_changes(names=names_to_keep)
348 def _process_importfrom_aliases(
349 self,
350 updated_node: cst.ImportFrom,
351 names: Iterable[cst.ImportAlias],
352 module_name: str,
353 ) -> Dict[str, Any]:
354 updates = {}
355 names_to_keep = []
356 objects_to_remove = self.unused_obj_imports[module_name]
357 for import_alias in names:
358 # Figure out if it is in our list of things to kill
359 for name, alias in objects_to_remove:
360 if (
361 name == import_alias.evaluated_name
362 and alias == import_alias.evaluated_alias
363 ):
364 break
365 else:
366 # This is a keeper, we don't have it on our list.
367 names_to_keep.append(import_alias)
368 continue
370 # Now that we know we want to remove this object, figure out if
371 # there are any live references to it.
372 if import_alias not in self._unused_imports:
373 names_to_keep.append(import_alias)
374 continue
376 # We are about to remove `import_alias`. Check if there are any
377 # trailing comments and reparent them to the previous import.
378 # We only do this in case there's a trailing comma, otherwise the
379 # entire import statement is going to be removed anyway.
380 comma = import_alias.comma
381 if isinstance(comma, cst.Comma):
382 if len(names_to_keep) != 0:
383 # there is a previous import alias
384 prev = names_to_keep[-1]
385 if isinstance(prev.comma, cst.Comma):
386 prev = prev.with_deep_changes(
387 prev.comma,
388 whitespace_after=_merge_whitespace_after(
389 prev.comma.whitespace_after,
390 comma.whitespace_after,
391 ),
392 )
393 else:
394 # The previous alias didn't have a trailing comma. This can
395 # occur if the alias was generated, instead of being parsed
396 # from source.
397 prev = prev.with_changes(comma=comma)
398 names_to_keep[-1] = prev
399 else:
400 # No previous import alias, need to attach comment to `ImportFrom`.
401 # We can only do this if there was a leftparen on the import
402 # statement. Otherwise there can't be any standalone comments
403 # anyway, so it's fine to skip this logic.
404 lpar = updated_node.lpar
405 if isinstance(lpar, cst.LeftParen):
406 updates["lpar"] = lpar.with_changes(
407 whitespace_after=_merge_whitespace_after(
408 lpar.whitespace_after,
409 comma.whitespace_after,
410 )
411 )
412 updates["names"] = names_to_keep
413 return updates
415 def leave_ImportFrom(
416 self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
417 ) -> Union[cst.ImportFrom, cst.RemovalSentinel]:
418 names = original_node.names
419 if isinstance(names, cst.ImportStar):
420 # This is a star import, so we won't remove it.
421 return updated_node
423 # Make sure we actually know the absolute module.
424 module_name = get_absolute_module_from_package_for_import(
425 self.context.full_package_name, updated_node
426 )
427 if module_name is None or module_name not in self.unused_obj_imports:
428 # This node isn't on our list of todos, so let's bail.
429 return updated_node
431 updates = self._process_importfrom_aliases(updated_node, names, module_name)
432 names_to_keep = updates["names"]
434 # no changes
435 if names_to_keep == names:
436 return updated_node
438 # Now, either remove this statement or remove the imports we are
439 # deleting from this statement.
440 if len(names_to_keep) == 0:
441 return cst.RemoveFromParent()
443 if names_to_keep[-1] != names[-1]:
444 # Remove trailing comma in order to not mess up import statements.
445 names_to_keep = [
446 *names_to_keep[:-1],
447 names_to_keep[-1].with_changes(comma=cst.MaybeSentinel.DEFAULT),
448 ]
449 updates["names"] = names_to_keep
450 return updated_node.with_changes(**updates)
453def _merge_whitespace_after(
454 left: cst.BaseParenthesizableWhitespace, right: cst.BaseParenthesizableWhitespace
455) -> cst.BaseParenthesizableWhitespace:
456 if not isinstance(right, cst.ParenthesizedWhitespace):
457 return left
458 if not isinstance(left, cst.ParenthesizedWhitespace):
459 return right
461 return left.with_changes(
462 empty_lines=tuple(
463 line for line in right.empty_lines if line.comment is not None
464 ),
465 )