Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/libcst/codemod/visitors/_gather_imports.py: 19%

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 Dict, List, Sequence, Set, Tuple, Union 

7 

8import libcst 

9from libcst.codemod._context import CodemodContext 

10from libcst.codemod._visitor import ContextAwareVisitor 

11from libcst.codemod.visitors._imports import ImportItem 

12from libcst.helpers import get_absolute_module_from_package_for_import 

13 

14 

15class GatherImportsVisitor(ContextAwareVisitor): 

16 """ 

17 Gathers all imports in a module and stores them as attributes on the instance. 

18 Intended to be instantiated and passed to a :class:`~libcst.Module` 

19 :meth:`~libcst.CSTNode.visit` method in order to gather up information about 

20 imports on a module. Note that this is not a substitute for scope analysis or 

21 qualified name support. Please see :ref:`libcst-scope-tutorial` for a more 

22 robust way of determining the qualified name and definition for an arbitrary 

23 node. 

24 

25 After visiting a module the following attributes will be populated: 

26 

27 module_imports 

28 A sequence of strings representing modules that were imported directly, such as 

29 in the case of ``import typing``. Each module directly imported but not aliased 

30 will be included here. 

31 object_mapping 

32 A mapping of strings to sequences of strings representing modules where we 

33 imported objects from, such as in the case of ``from typing import Optional``. 

34 Each from import that was not aliased will be included here, where the keys of 

35 the mapping are the module we are importing from, and the value is a 

36 sequence of objects we are importing from the module. 

37 module_aliases 

38 A mapping of strings representing modules that were imported and aliased, 

39 such as in the case of ``import typing as t``. Each module imported this 

40 way will be represented as a key in this mapping, and the value will be 

41 the local alias of the module. 

42 alias_mapping 

43 A mapping of strings to sequences of tuples representing modules where we 

44 imported objects from and aliased using ``as`` syntax, such as in the case 

45 of ``from typing import Optional as opt``. Each from import that was aliased 

46 will be included here, where the keys of the mapping are the module we are 

47 importing from, and the value is a tuple representing the original object 

48 name and the alias. 

49 all_imports 

50 A collection of all :class:`~libcst.Import` and :class:`~libcst.ImportFrom` 

51 statements that were encountered in the module. 

52 """ 

53 

54 def __init__(self, context: CodemodContext) -> None: 

55 super().__init__(context) 

56 # Track the available imports in this transform 

57 self.module_imports: Set[str] = set() 

58 self.object_mapping: Dict[str, Set[str]] = {} 

59 # Track the aliased imports in this transform 

60 self.module_aliases: Dict[str, str] = {} 

61 self.alias_mapping: Dict[str, List[Tuple[str, str]]] = {} 

62 # Track all of the imports found in this transform 

63 self.all_imports: List[Union[libcst.Import, libcst.ImportFrom]] = [] 

64 # Track the import for every symbol introduced into the module 

65 self.symbol_mapping: Dict[str, ImportItem] = {} 

66 

67 def visit_Import(self, node: libcst.Import) -> None: 

68 # Track this import statement for later analysis. 

69 self.all_imports.append(node) 

70 

71 for name in node.names: 

72 alias = name.evaluated_alias 

73 imp = ImportItem(name.evaluated_name, alias=alias) 

74 if alias is not None: 

75 # Track this as an aliased module 

76 self.module_aliases[name.evaluated_name] = alias 

77 self.symbol_mapping[alias] = imp 

78 else: 

79 # Get the module we're importing as a string. 

80 self.module_imports.add(name.evaluated_name) 

81 self.symbol_mapping[name.evaluated_name] = imp 

82 

83 def visit_ImportFrom(self, node: libcst.ImportFrom) -> None: 

84 # Track this import statement for later analysis. 

85 self.all_imports.append(node) 

86 

87 # Get the module we're importing as a string. 

88 module = get_absolute_module_from_package_for_import( 

89 self.context.full_package_name, node 

90 ) 

91 if module is None: 

92 # Can't get the absolute import from relative, so we can't 

93 # support this. 

94 return 

95 nodenames = node.names 

96 if isinstance(nodenames, libcst.ImportStar): 

97 # We cover everything, no need to bother tracking other things 

98 self.object_mapping[module] = set("*") 

99 return 

100 elif isinstance(nodenames, Sequence): 

101 # Get the list of imports we're aliasing in this import 

102 new_aliases = [ 

103 (ia.evaluated_name, ia.evaluated_alias) 

104 for ia in nodenames 

105 if ia.asname is not None 

106 ] 

107 if new_aliases: 

108 if module not in self.alias_mapping: 

109 self.alias_mapping[module] = [] 

110 # pyre-ignore We know that aliases are not None here. 

111 self.alias_mapping[module].extend(new_aliases) 

112 

113 # Get the list of imports we're importing in this import 

114 new_objects = {ia.evaluated_name for ia in nodenames if ia.asname is None} 

115 if new_objects: 

116 if module not in self.object_mapping: 

117 self.object_mapping[module] = set() 

118 

119 # Make sure that we don't add to a '*' module 

120 if "*" in self.object_mapping[module]: 

121 self.object_mapping[module] = set("*") 

122 return 

123 

124 self.object_mapping[module].update(new_objects) 

125 for ia in nodenames: 

126 imp = ImportItem( 

127 module, obj_name=ia.evaluated_name, alias=ia.evaluated_alias 

128 ) 

129 key = ia.evaluated_alias or ia.evaluated_name 

130 self.symbol_mapping[key] = imp