output.py 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689
  1. import copy
  2. import itertools
  3. from collections.abc import Iterable
  4. from functools import partial
  5. from typing import Any
  6. from isort.format import format_simplified
  7. from . import parse, sorting, wrap
  8. from .comments import add_to_line as with_comments
  9. from .identify import STATEMENT_DECLARATIONS
  10. from .settings import DEFAULT_CONFIG, Config
  11. # Ignore DeepSource cyclomatic complexity check for this function.
  12. # skipcq: PY-R1000
  13. def sorted_imports(
  14. parsed: parse.ParsedContent,
  15. config: Config = DEFAULT_CONFIG,
  16. extension: str = "py",
  17. import_type: str = "import",
  18. ) -> str:
  19. """Adds the imports back to the file.
  20. (at the index of the first import) sorted alphabetically and split between groups
  21. """
  22. if parsed.import_index == -1:
  23. return _output_as_string(parsed.lines_without_imports, parsed.line_separator)
  24. formatted_output: list[str] = parsed.lines_without_imports.copy()
  25. remove_imports = [format_simplified(removal) for removal in config.remove_imports]
  26. sections: Iterable[str] = itertools.chain(parsed.sections, config.forced_separate)
  27. if config.no_sections:
  28. parsed.imports["no_sections"] = {"straight": {}, "from": {}}
  29. base_sections: tuple[str, ...] = ()
  30. for section in sections:
  31. if section == "FUTURE":
  32. base_sections = ("FUTURE",)
  33. continue
  34. parsed.imports["no_sections"]["straight"].update(
  35. parsed.imports[section].get("straight", {})
  36. )
  37. parsed.imports["no_sections"]["from"].update(parsed.imports[section].get("from", {}))
  38. sections = (*base_sections, "no_sections")
  39. output: list[str] = []
  40. seen_headings: set[str] = set()
  41. pending_lines_before = False
  42. for section in sections:
  43. straight_modules = parsed.imports[section]["straight"]
  44. if not config.only_sections:
  45. straight_modules = sorting.sort(
  46. config,
  47. straight_modules,
  48. key=lambda key: sorting.module_key(
  49. key, config, section_name=section, straight_import=True
  50. ),
  51. reverse=config.reverse_sort,
  52. )
  53. from_modules = parsed.imports[section]["from"]
  54. if not config.only_sections:
  55. from_modules = sorting.sort(
  56. config,
  57. from_modules,
  58. key=lambda key: sorting.module_key(key, config, section_name=section),
  59. reverse=config.reverse_sort,
  60. )
  61. if config.star_first:
  62. star_modules = []
  63. other_modules = []
  64. for module in from_modules:
  65. if "*" in parsed.imports[section]["from"][module]:
  66. star_modules.append(module)
  67. else:
  68. other_modules.append(module)
  69. from_modules = star_modules + other_modules
  70. straight_imports = _with_straight_imports(
  71. parsed, config, straight_modules, section, remove_imports, import_type
  72. )
  73. from_imports = _with_from_imports(
  74. parsed, config, from_modules, section, remove_imports, import_type
  75. )
  76. lines_between = [""] * (
  77. config.lines_between_types if from_modules and straight_modules else 0
  78. )
  79. if config.from_first or section == "FUTURE":
  80. section_output = from_imports + lines_between + straight_imports
  81. else:
  82. section_output = straight_imports + lines_between + from_imports
  83. if config.force_sort_within_sections:
  84. # collapse comments
  85. comments_above = []
  86. new_section_output: list[str] = []
  87. for line in section_output:
  88. if not line:
  89. continue
  90. if line.startswith("#"):
  91. comments_above.append(line)
  92. elif comments_above:
  93. new_section_output.append(_LineWithComments(line, comments_above))
  94. comments_above = []
  95. else:
  96. new_section_output.append(line)
  97. # only_sections options is not imposed if force_sort_within_sections is True
  98. new_section_output = sorting.sort(
  99. config,
  100. new_section_output,
  101. key=partial(sorting.section_key, config=config),
  102. reverse=config.reverse_sort,
  103. )
  104. # uncollapse comments
  105. section_output = []
  106. for line in new_section_output:
  107. comments = getattr(line, "comments", ())
  108. if comments:
  109. section_output.extend(comments)
  110. section_output.append(str(line))
  111. section_name = section
  112. no_lines_before = section_name in config.no_lines_before
  113. if section_output:
  114. if section_name in parsed.place_imports:
  115. parsed.place_imports[section_name] = section_output
  116. continue
  117. section_title = config.import_headings.get(section_name.lower(), "")
  118. if section_title and section_title not in seen_headings:
  119. if config.dedup_headings:
  120. seen_headings.add(section_title)
  121. section_comment = f"# {section_title}"
  122. if section_comment not in parsed.lines_without_imports[0:1]: # pragma: no branch
  123. section_output.insert(0, section_comment)
  124. section_footer = config.import_footers.get(section_name.lower(), "")
  125. if section_footer and section_footer not in seen_headings:
  126. if config.dedup_headings:
  127. seen_headings.add(section_footer)
  128. section_comment_end = f"# {section_footer}"
  129. if (
  130. section_comment_end not in parsed.lines_without_imports[-1:]
  131. ): # pragma: no branch
  132. section_output.append("") # Empty line for black compatibility
  133. section_output.append(section_comment_end)
  134. if pending_lines_before or not no_lines_before:
  135. output += [""] * config.lines_between_sections
  136. output += section_output
  137. pending_lines_before = False
  138. else:
  139. pending_lines_before = pending_lines_before or not no_lines_before
  140. if config.ensure_newline_before_comments:
  141. output = _ensure_newline_before_comment(output)
  142. while output and output[-1].strip() == "":
  143. output.pop() # pragma: no cover
  144. while output and output[0].strip() == "":
  145. output.pop(0)
  146. if config.formatting_function:
  147. output = config.formatting_function(
  148. parsed.line_separator.join(output), extension, config
  149. ).splitlines()
  150. output_at = 0
  151. if parsed.import_index < parsed.original_line_count:
  152. output_at = parsed.import_index
  153. formatted_output[output_at:0] = output
  154. if output:
  155. imports_tail = output_at + len(output)
  156. while [
  157. character.strip() for character in formatted_output[imports_tail : imports_tail + 1]
  158. ] == [""]:
  159. formatted_output.pop(imports_tail)
  160. if config.lines_before_imports != -1:
  161. lines_before_imports = config.lines_before_imports
  162. if config.profile == "black" and extension == "pyi": # special case for black
  163. lines_before_imports = 1
  164. formatted_output[:0] = ["" for line in range(lines_before_imports)]
  165. imports_tail += lines_before_imports
  166. if len(formatted_output) > imports_tail:
  167. next_construct = ""
  168. tail = formatted_output[imports_tail:]
  169. for index, line in enumerate(tail): # pragma: no branch
  170. should_skip, in_quote, *_ = parse.skip_line(
  171. line,
  172. in_quote="",
  173. index=len(formatted_output),
  174. section_comments=config.section_comments,
  175. needs_import=False,
  176. )
  177. if not should_skip and line.strip():
  178. if (
  179. line.strip().startswith("#")
  180. and len(tail) > (index + 1)
  181. and tail[index + 1].strip()
  182. ):
  183. continue
  184. next_construct = line
  185. break
  186. if in_quote: # pragma: no branch
  187. next_construct = line
  188. break
  189. if config.lines_after_imports != -1:
  190. lines_after_imports = config.lines_after_imports
  191. if config.profile == "black" and extension == "pyi": # special case for black
  192. lines_after_imports = 1
  193. formatted_output[imports_tail:0] = ["" for line in range(lines_after_imports)]
  194. elif extension != "pyi" and next_construct.startswith(STATEMENT_DECLARATIONS):
  195. formatted_output[imports_tail:0] = ["", ""]
  196. else:
  197. formatted_output[imports_tail:0] = [""]
  198. if parsed.place_imports:
  199. new_out_lines = []
  200. for index, line in enumerate(formatted_output):
  201. new_out_lines.append(line)
  202. if line in parsed.import_placements:
  203. new_out_lines.extend(parsed.place_imports[parsed.import_placements[line]])
  204. if (
  205. len(formatted_output) <= (index + 1)
  206. or formatted_output[index + 1].strip() != ""
  207. ):
  208. new_out_lines.append("")
  209. formatted_output = new_out_lines
  210. return _output_as_string(formatted_output, parsed.line_separator)
  211. # Ignore DeepSource cyclomatic complexity check for this function. It was
  212. # already complex when this check was enabled.
  213. # skipcq: PY-R1000
  214. def _with_from_imports(
  215. parsed: parse.ParsedContent,
  216. config: Config,
  217. from_modules: Iterable[str],
  218. section: str,
  219. remove_imports: list[str],
  220. import_type: str,
  221. ) -> list[str]:
  222. output: list[str] = []
  223. for module in from_modules:
  224. if module in remove_imports:
  225. continue
  226. import_start = f"from {module} {import_type} "
  227. from_imports = list(parsed.imports[section]["from"][module])
  228. if (
  229. not config.no_inline_sort
  230. or (config.force_single_line and module not in config.single_line_exclusions)
  231. ) and not config.only_sections:
  232. from_imports = sorting.sort(
  233. config,
  234. from_imports,
  235. key=lambda key: sorting.module_key(
  236. key,
  237. config,
  238. True,
  239. config.force_alphabetical_sort_within_sections,
  240. section_name=section,
  241. ),
  242. reverse=config.reverse_sort,
  243. )
  244. if remove_imports:
  245. from_imports = [
  246. line for line in from_imports if f"{module}.{line}" not in remove_imports
  247. ]
  248. sub_modules = [f"{module}.{from_import}" for from_import in from_imports]
  249. as_imports = {
  250. from_import: [
  251. f"{from_import} as {as_module}" for as_module in parsed.as_map["from"][sub_module]
  252. ]
  253. for from_import, sub_module in zip(from_imports, sub_modules, strict=False)
  254. if sub_module in parsed.as_map["from"]
  255. }
  256. if config.combine_as_imports and not ("*" in from_imports and config.combine_star):
  257. if not config.no_inline_sort:
  258. for as_import in as_imports:
  259. if not config.only_sections:
  260. as_imports[as_import] = sorting.sort(config, as_imports[as_import])
  261. for from_import in copy.copy(from_imports):
  262. if from_import in as_imports:
  263. idx = from_imports.index(from_import)
  264. if parsed.imports[section]["from"][module][from_import]:
  265. from_imports[(idx + 1) : (idx + 1)] = as_imports.pop(from_import)
  266. else:
  267. from_imports[idx : (idx + 1)] = as_imports.pop(from_import)
  268. only_show_as_imports = False
  269. comments = parsed.categorized_comments["from"].pop(module, ())
  270. above_comments = parsed.categorized_comments["above"]["from"].pop(module, None)
  271. while from_imports:
  272. if above_comments:
  273. output.extend(above_comments)
  274. above_comments = None
  275. if "*" in from_imports and config.combine_star:
  276. import_statement = wrap.line(
  277. with_comments(
  278. _with_star_comments(parsed, module, list(comments or ())),
  279. f"{import_start}*",
  280. removed=config.ignore_comments,
  281. comment_prefix=config.comment_prefix,
  282. ),
  283. parsed.line_separator,
  284. config,
  285. )
  286. from_imports = [
  287. from_import for from_import in from_imports if from_import in as_imports
  288. ]
  289. only_show_as_imports = True
  290. elif config.force_single_line and module not in config.single_line_exclusions:
  291. import_statement = ""
  292. while from_imports:
  293. from_import = from_imports.pop(0)
  294. single_import_line = with_comments(
  295. comments,
  296. import_start + from_import,
  297. removed=config.ignore_comments,
  298. comment_prefix=config.comment_prefix,
  299. )
  300. comment = (
  301. parsed.categorized_comments["nested"].get(module, {}).pop(from_import, None)
  302. )
  303. if comment:
  304. single_import_line += (
  305. f"{(comments and ';') or config.comment_prefix} {comment}"
  306. )
  307. if from_import in as_imports:
  308. if (
  309. parsed.imports[section]["from"][module][from_import]
  310. and not only_show_as_imports
  311. ):
  312. output.append(
  313. wrap.line(single_import_line, parsed.line_separator, config)
  314. )
  315. from_comments = parsed.categorized_comments["straight"].get(
  316. f"{module}.{from_import}"
  317. )
  318. if not config.only_sections:
  319. output.extend(
  320. with_comments(
  321. from_comments,
  322. wrap.line(
  323. import_start + as_import, parsed.line_separator, config
  324. ),
  325. removed=config.ignore_comments,
  326. comment_prefix=config.comment_prefix,
  327. )
  328. for as_import in sorting.sort(config, as_imports[from_import])
  329. )
  330. else:
  331. output.extend(
  332. with_comments(
  333. from_comments,
  334. wrap.line(
  335. import_start + as_import, parsed.line_separator, config
  336. ),
  337. removed=config.ignore_comments,
  338. comment_prefix=config.comment_prefix,
  339. )
  340. for as_import in as_imports[from_import]
  341. )
  342. else:
  343. output.append(wrap.line(single_import_line, parsed.line_separator, config))
  344. comments = None
  345. else:
  346. while from_imports and from_imports[0] in as_imports:
  347. from_import = from_imports.pop(0)
  348. if not config.only_sections:
  349. as_imports[from_import] = sorting.sort(config, as_imports[from_import])
  350. from_comments = (
  351. parsed.categorized_comments["straight"].get(f"{module}.{from_import}") or []
  352. )
  353. if (
  354. parsed.imports[section]["from"][module][from_import]
  355. and not only_show_as_imports
  356. ):
  357. specific_comment = (
  358. parsed.categorized_comments["nested"]
  359. .get(module, {})
  360. .pop(from_import, None)
  361. )
  362. if specific_comment:
  363. from_comments.append(specific_comment)
  364. output.append(
  365. wrap.line(
  366. with_comments(
  367. from_comments,
  368. import_start + from_import,
  369. removed=config.ignore_comments,
  370. comment_prefix=config.comment_prefix,
  371. ),
  372. parsed.line_separator,
  373. config,
  374. )
  375. )
  376. from_comments = []
  377. for as_import in as_imports[from_import]:
  378. specific_comment = (
  379. parsed.categorized_comments["nested"]
  380. .get(module, {})
  381. .pop(as_import, None)
  382. )
  383. if specific_comment:
  384. from_comments.append(specific_comment)
  385. output.append(
  386. wrap.line(
  387. with_comments(
  388. from_comments,
  389. import_start + as_import,
  390. removed=config.ignore_comments,
  391. comment_prefix=config.comment_prefix,
  392. ),
  393. parsed.line_separator,
  394. config,
  395. )
  396. )
  397. from_comments = []
  398. if "*" in from_imports:
  399. output.append(
  400. with_comments(
  401. _with_star_comments(parsed, module, []),
  402. f"{import_start}*",
  403. removed=config.ignore_comments,
  404. comment_prefix=config.comment_prefix,
  405. )
  406. )
  407. from_imports.remove("*")
  408. for from_import in copy.copy(from_imports):
  409. comment = (
  410. parsed.categorized_comments["nested"].get(module, {}).pop(from_import, None)
  411. )
  412. if comment:
  413. # If the comment is a noqa and hanging indent wrapping is used,
  414. # keep the name in the main list and hoist the comment to the statement.
  415. if (
  416. comment.lower().startswith("noqa")
  417. and config.multi_line_output == wrap.Modes.HANGING_INDENT # type: ignore[attr-defined] # noqa: E501
  418. ):
  419. comments = list(comments) if comments else []
  420. comments.append(comment)
  421. continue
  422. from_imports.remove(from_import)
  423. if from_imports:
  424. use_comments = []
  425. else:
  426. use_comments = comments
  427. comments = None
  428. single_import_line = with_comments(
  429. use_comments,
  430. import_start + from_import,
  431. removed=config.ignore_comments,
  432. comment_prefix=config.comment_prefix,
  433. )
  434. single_import_line += (
  435. f"{(use_comments and ';') or config.comment_prefix} {comment}"
  436. )
  437. output.append(wrap.line(single_import_line, parsed.line_separator, config))
  438. from_import_section = []
  439. while from_imports and (
  440. from_imports[0] not in as_imports
  441. or (
  442. config.combine_as_imports
  443. and parsed.imports[section]["from"][module][from_import]
  444. )
  445. ):
  446. from_import_section.append(from_imports.pop(0))
  447. if config.combine_as_imports:
  448. comments = (comments or []) + list(
  449. parsed.categorized_comments["from"].pop(f"{module}.__combined_as__", ())
  450. )
  451. import_statement = with_comments(
  452. comments,
  453. import_start + (", ").join(from_import_section),
  454. removed=config.ignore_comments,
  455. comment_prefix=config.comment_prefix,
  456. )
  457. if not from_import_section:
  458. import_statement = ""
  459. do_multiline_reformat = False
  460. force_grid_wrap = config.force_grid_wrap
  461. if force_grid_wrap and len(from_import_section) >= force_grid_wrap:
  462. do_multiline_reformat = True
  463. if len(import_statement) > config.line_length and len(from_import_section) > 1:
  464. do_multiline_reformat = True
  465. # If line too long AND have imports AND we are
  466. # NOT using GRID or VERTICAL wrap modes
  467. if (
  468. len(import_statement) > config.line_length
  469. and len(from_import_section) > 0
  470. and config.multi_line_output not in (wrap.Modes.GRID, wrap.Modes.VERTICAL) # type: ignore # noqa: E501
  471. ):
  472. do_multiline_reformat = True
  473. if (
  474. import_statement
  475. and config.split_on_trailing_comma
  476. and module in parsed.trailing_commas
  477. ):
  478. import_statement = wrap.import_statement(
  479. import_start=import_start,
  480. from_imports=from_import_section,
  481. comments=comments,
  482. line_separator=parsed.line_separator,
  483. config=config,
  484. explode=True,
  485. )
  486. elif do_multiline_reformat:
  487. import_statement = wrap.import_statement(
  488. import_start=import_start,
  489. from_imports=from_import_section,
  490. comments=comments,
  491. line_separator=parsed.line_separator,
  492. config=config,
  493. )
  494. if config.multi_line_output == wrap.Modes.GRID: # type: ignore
  495. other_import_statement = wrap.import_statement(
  496. import_start=import_start,
  497. from_imports=from_import_section,
  498. comments=comments,
  499. line_separator=parsed.line_separator,
  500. config=config,
  501. multi_line_output=wrap.Modes.VERTICAL_GRID, # type: ignore
  502. )
  503. if (
  504. max(
  505. len(import_line)
  506. for import_line in import_statement.split(parsed.line_separator)
  507. )
  508. > config.line_length
  509. ):
  510. import_statement = other_import_statement
  511. elif len(import_statement) > config.line_length:
  512. import_statement = wrap.line(import_statement, parsed.line_separator, config)
  513. if import_statement:
  514. output.append(import_statement)
  515. return output
  516. def _with_straight_imports(
  517. parsed: parse.ParsedContent,
  518. config: Config,
  519. straight_modules: Iterable[str],
  520. section: str,
  521. remove_imports: list[str],
  522. import_type: str,
  523. ) -> list[str]:
  524. output: list[str] = []
  525. as_imports = any(module in parsed.as_map["straight"] for module in straight_modules)
  526. # combine_straight_imports only works for bare imports, 'as' imports not included
  527. if config.combine_straight_imports and not as_imports:
  528. if not straight_modules:
  529. return []
  530. above_comments: list[str] = []
  531. inline_comments: list[str] = []
  532. for module in straight_modules:
  533. if module in parsed.categorized_comments["above"]["straight"]:
  534. above_comments.extend(parsed.categorized_comments["above"]["straight"].pop(module))
  535. if module in parsed.categorized_comments["straight"]:
  536. inline_comments.extend(parsed.categorized_comments["straight"][module])
  537. combined_straight_imports = ", ".join(straight_modules)
  538. if inline_comments:
  539. combined_inline_comments = " ".join(inline_comments)
  540. else:
  541. combined_inline_comments = ""
  542. output.extend(above_comments)
  543. if combined_inline_comments:
  544. output.append(
  545. f"{import_type} {combined_straight_imports} # {combined_inline_comments}"
  546. )
  547. else:
  548. output.append(f"{import_type} {combined_straight_imports}")
  549. return output
  550. for module in straight_modules:
  551. if module in remove_imports:
  552. continue
  553. import_definition = []
  554. if module in parsed.as_map["straight"]:
  555. if parsed.imports[section]["straight"][module]:
  556. import_definition.append((f"{import_type} {module}", module))
  557. import_definition.extend(
  558. (f"{import_type} {module} as {as_import}", f"{module} as {as_import}")
  559. for as_import in parsed.as_map["straight"][module]
  560. )
  561. else:
  562. import_definition.append((f"{import_type} {module}", module))
  563. comments_above = parsed.categorized_comments["above"]["straight"].pop(module, None)
  564. if comments_above:
  565. output.extend(comments_above)
  566. output.extend(
  567. with_comments(
  568. parsed.categorized_comments["straight"].get(imodule),
  569. idef,
  570. removed=config.ignore_comments,
  571. comment_prefix=config.comment_prefix,
  572. )
  573. for idef, imodule in import_definition
  574. )
  575. return output
  576. def _output_as_string(lines: list[str], line_separator: str) -> str:
  577. return line_separator.join(_normalize_empty_lines(lines))
  578. def _normalize_empty_lines(lines: list[str]) -> list[str]:
  579. while lines and lines[-1].strip() == "":
  580. lines.pop(-1)
  581. lines.append("")
  582. return lines
  583. class _LineWithComments(str):
  584. comments: list[str]
  585. def __new__(
  586. cls: type["_LineWithComments"], value: Any, comments: list[str]
  587. ) -> "_LineWithComments":
  588. instance = super().__new__(cls, value)
  589. instance.comments = comments
  590. return instance
  591. def _ensure_newline_before_comment(output: list[str]) -> list[str]:
  592. new_output: list[str] = []
  593. def is_comment(line: str | None) -> bool:
  594. return line.startswith("#") if line else False
  595. for line, prev_line in zip(output, [None, *output], strict=False):
  596. if is_comment(line) and prev_line != "" and not is_comment(prev_line):
  597. new_output.append("")
  598. new_output.append(line)
  599. return new_output
  600. def _with_star_comments(parsed: parse.ParsedContent, module: str, comments: list[str]) -> list[str]:
  601. star_comment = parsed.categorized_comments["nested"].get(module, {}).pop("*", None)
  602. if star_comment:
  603. return [*comments, star_comment]
  604. return comments