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