identify.py 8.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208
  1. """Fast stream based import identification.
  2. Eventually this will likely replace parse.py
  3. """
  4. from collections.abc import Iterator
  5. from functools import partial
  6. from pathlib import Path
  7. from typing import NamedTuple, TextIO
  8. from isort.parse import normalize_line, skip_line, strip_syntax
  9. from .comments import parse as parse_comments
  10. from .settings import DEFAULT_CONFIG, Config
  11. STATEMENT_DECLARATIONS: tuple[str, ...] = ("def ", "cdef ", "cpdef ", "class ", "@", "async def")
  12. class Import(NamedTuple):
  13. line_number: int
  14. indented: bool
  15. module: str
  16. attribute: str | None = None
  17. alias: str | None = None
  18. cimport: bool = False
  19. file_path: Path | None = None
  20. def statement(self) -> str:
  21. import_cmd = "cimport" if self.cimport else "import"
  22. if self.attribute:
  23. import_string = f"from {self.module} {import_cmd} {self.attribute}"
  24. else:
  25. import_string = f"{import_cmd} {self.module}"
  26. if self.alias:
  27. import_string += f" as {self.alias}"
  28. return import_string
  29. def __str__(self) -> str:
  30. return (
  31. f"{self.file_path or ''}:{self.line_number} "
  32. f"{'indented ' if self.indented else ''}{self.statement()}"
  33. )
  34. def imports(
  35. input_stream: TextIO,
  36. config: Config = DEFAULT_CONFIG,
  37. file_path: Path | None = None,
  38. top_only: bool = False,
  39. ) -> Iterator[Import]:
  40. """Parses a python file taking out and categorizing imports."""
  41. in_quote = ""
  42. indexed_input = enumerate(input_stream)
  43. for index, raw_line in indexed_input:
  44. (skipping_line, in_quote) = skip_line(
  45. raw_line, in_quote=in_quote, index=index, section_comments=config.section_comments
  46. )
  47. if top_only and not in_quote and raw_line.startswith(STATEMENT_DECLARATIONS):
  48. break
  49. if skipping_line:
  50. continue
  51. stripped_line = raw_line.strip().split("#")[0]
  52. if stripped_line.startswith(("raise", "yield")):
  53. if stripped_line == "yield":
  54. while not stripped_line or stripped_line == "yield":
  55. try:
  56. index, next_line = next(indexed_input)
  57. except StopIteration:
  58. break
  59. stripped_line = next_line.strip().split("#")[0]
  60. while stripped_line.endswith("\\"):
  61. try:
  62. index, next_line = next(indexed_input)
  63. except StopIteration:
  64. break
  65. stripped_line = next_line.strip().split("#")[0]
  66. continue # pragma: no cover
  67. line, *end_of_line_comment = raw_line.split("#", 1)
  68. statements = [line.strip() for line in line.split(";")]
  69. if end_of_line_comment:
  70. statements[-1] = f"{statements[-1]}#{end_of_line_comment[0]}"
  71. for statement in statements:
  72. line, _raw_line = normalize_line(statement)
  73. if line.startswith(("import ", "cimport ")):
  74. type_of_import = "straight"
  75. elif line.startswith("from "):
  76. type_of_import = "from"
  77. else:
  78. continue # pragma: no cover
  79. import_string, _ = parse_comments(line)
  80. normalized_import_string = (
  81. import_string.replace("import(", "import (").replace("\\", " ").replace("\n", " ")
  82. )
  83. cimports: bool = (
  84. " cimport " in normalized_import_string
  85. or normalized_import_string.startswith("cimport")
  86. )
  87. identified_import = partial(
  88. Import,
  89. index + 1, # line numbers use 1 based indexing
  90. raw_line.startswith((" ", "\t")),
  91. cimport=cimports,
  92. file_path=file_path,
  93. )
  94. if "(" in line.split("#", 1)[0]:
  95. while not line.split("#")[0].strip().endswith(")"):
  96. try:
  97. index, next_line = next(indexed_input)
  98. except StopIteration:
  99. break
  100. line, _ = parse_comments(next_line)
  101. import_string += "\n" + line
  102. else:
  103. while line.strip().endswith("\\"):
  104. try:
  105. index, next_line = next(indexed_input)
  106. except StopIteration:
  107. break
  108. line, _ = parse_comments(next_line)
  109. # Still need to check for parentheses after an escaped line
  110. if "(" in line.split("#")[0] and ")" not in line.split("#")[0]:
  111. import_string += "\n" + line
  112. while not line.split("#")[0].strip().endswith(")"):
  113. try:
  114. index, next_line = next(indexed_input)
  115. except StopIteration:
  116. break
  117. line, _ = parse_comments(next_line)
  118. import_string += "\n" + line
  119. else:
  120. if import_string.strip().endswith(
  121. (" import", " cimport")
  122. ) or line.strip().startswith(("import ", "cimport ")):
  123. import_string += "\n" + line
  124. else:
  125. import_string = (
  126. import_string.rstrip().rstrip("\\") + " " + line.lstrip()
  127. )
  128. if type_of_import == "from":
  129. import_string = (
  130. import_string.replace("import(", "import (")
  131. .replace("\\", " ")
  132. .replace("\n", " ")
  133. )
  134. parts = import_string.split(" cimport " if cimports else " import ")
  135. from_import = parts[0].split(" ")
  136. import_string = (" cimport " if cimports else " import ").join(
  137. [from_import[0] + " " + "".join(from_import[1:]), *parts[1:]]
  138. )
  139. just_imports = [
  140. item.replace("{|", "{ ").replace("|}", " }")
  141. for item in strip_syntax(import_string).split()
  142. ]
  143. direct_imports = just_imports[1:]
  144. top_level_module = ""
  145. if "as" in just_imports and (just_imports.index("as") + 1) < len(just_imports):
  146. while "as" in just_imports:
  147. attribute = None
  148. as_index = just_imports.index("as")
  149. if type_of_import == "from":
  150. attribute = just_imports[as_index - 1]
  151. top_level_module = just_imports[0]
  152. module = top_level_module + "." + attribute
  153. alias = just_imports[as_index + 1]
  154. direct_imports.remove(attribute)
  155. direct_imports.remove(alias)
  156. direct_imports.remove("as")
  157. just_imports[1:] = direct_imports
  158. if attribute == alias and config.remove_redundant_aliases:
  159. yield identified_import(top_level_module, attribute)
  160. else:
  161. yield identified_import(top_level_module, attribute, alias=alias)
  162. else:
  163. module = just_imports[as_index - 1]
  164. alias = just_imports[as_index + 1]
  165. just_imports.remove(alias)
  166. just_imports.remove("as")
  167. just_imports.remove(module)
  168. if module == alias and config.remove_redundant_aliases:
  169. yield identified_import(module)
  170. else:
  171. yield identified_import(module, alias=alias)
  172. if just_imports:
  173. if type_of_import == "from":
  174. module = just_imports.pop(0)
  175. for attribute in just_imports:
  176. yield identified_import(module, attribute)
  177. else:
  178. for module in just_imports:
  179. yield identified_import(module)