Coverage for circular_deps/parsers/python.py: 64%

1 from __future__ import annotations

2

3 import tree_sitter_python as tspython

4 from tree_sitter import Language

5 from tree_sitter import Parser

6

7 from circular_deps.models import ImportInfo

8 from circular_deps.parsers.base import ImportParser

9

10 PY_LANGUAGE = Language(tspython.language())

11

12 _IMPORT_NODE_TYPES = {"import_statement", "import_from_statement"}

13 _DYNAMIC_IMPORTS = {"__import__", "importlib.import_module"}

14

15

16 def parse(code):

17 parser_obj = Parser(PY_LANGUAGE)

18 tree = parser_obj.parse(code.encode())

19 return tree.root_node

20

21

22 class PythonImportParser(ImportParser):

23 def extract_imports(self, source, filepath):

24 root = parse(source)

25 imports = []

26

27 for node in root.children:

28 if node.type == "import_statement":

29 imports.extend(self._process_import_statement(node))

30 elif node.type == "import_from_statement":

31 imports.extend(self._process_import_from_statement(node))

32

33 return self._filter_dynamic(imports, root)

34

35 def _process_import_statement(self, node):

36 imports = []

37 has_aliased = False

38

39 for child in node.children:

40 if child.type == "aliased_import":

41 has_aliased = True

42 name = ""

43 for nc in child.children:

44 if nc.type == "dotted_name":

45 name = nc.text.decode()

46 elif nc.type == "identifier":

47 name = nc.text.decode()

48 if name:

49 imports.append(

50 ImportInfo(

51 raw_module=name,

52 resolved_path=None,

53 import_type="absolute",

54 line=node.start_point[0] + 1,

55 is_dynamic=False,

56 node_type="import_statement",

57 )

58 )

59 elif not has_aliased and child.type == "dotted_name":

60 name = child.text.decode()

61 if name:

62 imports.append(

63 ImportInfo(

64 raw_module=name,

65 resolved_path=None,

66 import_type="absolute",

67 line=node.start_point[0] + 1,

68 is_dynamic=False,

69 node_type="import_statement",

70 )

71 )

72 return imports

73

74 def _process_import_from_statement(self, node):

75 imports = []

76 module_name = ""

77 is_relative = False

78 found_import = False

79

80 for child in node.children:

81 if child.type == "import":

82 found_import = True

83 break

84 if not found_import:

85 if child.type == "relative_import":

86 is_relative = True

87 module_name = child.text.decode()

88 elif child.type == "dotted_name" and not module_name:

89 module_name = child.text.decode()

90 elif child.type == "module_name":

91 for mc in child.children:

92 if mc.type in ("dotted_name", "relative_import"):

93 module_name = mc.text.decode()

94

95 if module_name:

96 imports.append(

97 ImportInfo(

98 raw_module=module_name,

99 resolved_path=None,

100 import_type="relative" if is_relative else "absolute",

101 line=node.start_point[0] + 1,

102 is_dynamic=False,

103 node_type="import_from_statement",

104 )

105 )

106 return imports

107

108 def _filter_dynamic(self, imports, root):

109 for node in root.children:

110 if node.type == "expression_statement":

111 for expr in node.children:

112 if expr.type == "call":

113 func_name = self._get_function_name(expr)

114 if func_name in _DYNAMIC_IMPORTS:

115 for imp in imports:

116 if imp.line == node.start_point[0] + 1:

117 imp.is_dynamic = True

118

119 return [imp for imp in imports if not imp.is_dynamic]

120

121 def _get_function_name(self, call_node):

122 for child in call_node.children:

123 if child.type == "attribute":

124 for gc in child.children:

125 if gc.type == "identifier":

126 name = gc.text.decode()

127 for nc in child.children:

128 if nc.type == "identifier" and nc != gc:

129 parent_name = nc.text.decode()

130 return f"{parent_name}.{name}"

131 elif child.type == "identifier":

132 return child.text.decode()

133 return ""