wrap_modes.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375
  1. """Defines all wrap modes that can be used when outputting formatted imports"""
  2. import enum
  3. from collections.abc import Callable
  4. from inspect import signature
  5. from typing import Any
  6. import isort.comments
  7. _wrap_modes: dict[str, Callable[..., str]] = {}
  8. def from_string(value: str) -> "WrapModes":
  9. return getattr(WrapModes, str(value), None) or WrapModes(int(value))
  10. def formatter_from_string(name: str) -> Callable[..., str]:
  11. return _wrap_modes.get(name.upper(), grid)
  12. def _wrap_mode_interface(
  13. statement: str,
  14. imports: list[str],
  15. white_space: str,
  16. indent: str,
  17. line_length: int,
  18. comments: list[str],
  19. line_separator: str,
  20. comment_prefix: str,
  21. include_trailing_comma: bool,
  22. remove_comments: bool,
  23. ) -> str:
  24. """Defines the common interface used by all wrap mode functions"""
  25. return ""
  26. def _wrap_mode(function: Callable[..., str]) -> Callable[..., str]:
  27. """Registers an individual wrap mode. Function name and order are significant and used for
  28. creating enum.
  29. """
  30. _wrap_modes[function.__name__.upper()] = function
  31. function.__signature__ = signature(_wrap_mode_interface) # type: ignore
  32. function.__annotations__ = _wrap_mode_interface.__annotations__
  33. return function
  34. @_wrap_mode
  35. def grid(**interface: Any) -> str:
  36. if not interface["imports"]:
  37. return ""
  38. interface["statement"] += "(" + interface["imports"].pop(0)
  39. while interface["imports"]:
  40. next_import = interface["imports"].pop(0)
  41. next_statement = isort.comments.add_to_line(
  42. interface["comments"],
  43. interface["statement"] + ", " + next_import,
  44. removed=interface["remove_comments"],
  45. comment_prefix=interface["comment_prefix"],
  46. )
  47. if (
  48. len(next_statement.split(interface["line_separator"])[-1]) + 1
  49. > interface["line_length"]
  50. ):
  51. lines = [f"{interface['white_space']}{next_import.split(' ')[0]}"]
  52. for part in next_import.split(" ")[1:]:
  53. new_line = f"{lines[-1]} {part}"
  54. if len(new_line) + 1 > interface["line_length"]:
  55. lines.append(f"{interface['white_space']}{part}")
  56. else:
  57. lines[-1] = new_line
  58. next_import = interface["line_separator"].join(lines)
  59. interface["statement"] = (
  60. isort.comments.add_to_line(
  61. interface["comments"],
  62. f"{interface['statement']},",
  63. removed=interface["remove_comments"],
  64. comment_prefix=interface["comment_prefix"],
  65. )
  66. + f"{interface['line_separator']}{next_import}"
  67. )
  68. interface["comments"] = []
  69. else:
  70. interface["statement"] += ", " + next_import
  71. return f"{interface['statement']}{',' if interface['include_trailing_comma'] else ''})"
  72. @_wrap_mode
  73. def vertical(**interface: Any) -> str:
  74. if not interface["imports"]:
  75. return ""
  76. first_import = (
  77. isort.comments.add_to_line(
  78. interface["comments"],
  79. interface["imports"].pop(0) + ",",
  80. removed=interface["remove_comments"],
  81. comment_prefix=interface["comment_prefix"],
  82. )
  83. + interface["line_separator"]
  84. + interface["white_space"]
  85. )
  86. _imports = ("," + interface["line_separator"] + interface["white_space"]).join(
  87. interface["imports"]
  88. )
  89. _comma_maybe = "," if interface["include_trailing_comma"] else ""
  90. return f"{interface['statement']}({first_import}{_imports}{_comma_maybe})"
  91. def _hanging_indent_end_line(line: str) -> str:
  92. if not line.endswith(" "):
  93. line += " "
  94. return line + "\\"
  95. @_wrap_mode
  96. def hanging_indent(**interface: Any) -> str:
  97. if not interface["imports"]:
  98. return ""
  99. line_length_limit = interface["line_length"] - 3
  100. next_import = interface["imports"].pop(0)
  101. next_statement = interface["statement"] + next_import
  102. # Check for first import
  103. if len(next_statement) > line_length_limit:
  104. next_statement = (
  105. _hanging_indent_end_line(interface["statement"])
  106. + interface["line_separator"]
  107. + interface["indent"]
  108. + next_import
  109. )
  110. interface["statement"] = next_statement
  111. while interface["imports"]:
  112. next_import = interface["imports"].pop(0)
  113. next_statement = interface["statement"] + ", " + next_import
  114. if len(next_statement.split(interface["line_separator"])[-1]) > line_length_limit:
  115. next_statement = (
  116. _hanging_indent_end_line(interface["statement"] + ",")
  117. + f"{interface['line_separator']}{interface['indent']}{next_import}"
  118. )
  119. interface["statement"] = next_statement
  120. if interface["comments"]:
  121. statement_with_comments = isort.comments.add_to_line(
  122. interface["comments"],
  123. interface["statement"],
  124. removed=interface["remove_comments"],
  125. comment_prefix=interface["comment_prefix"],
  126. )
  127. if len(statement_with_comments.split(interface["line_separator"])[-1]) <= (
  128. line_length_limit + 2
  129. ):
  130. return statement_with_comments
  131. return (
  132. _hanging_indent_end_line(interface["statement"])
  133. + str(interface["line_separator"])
  134. + isort.comments.add_to_line(
  135. interface["comments"],
  136. interface["indent"],
  137. removed=interface["remove_comments"],
  138. comment_prefix=interface["comment_prefix"].lstrip(),
  139. )
  140. )
  141. return str(interface["statement"])
  142. @_wrap_mode
  143. def vertical_hanging_indent(**interface: Any) -> str:
  144. _line_with_comments = isort.comments.add_to_line(
  145. interface["comments"],
  146. "",
  147. removed=interface["remove_comments"],
  148. comment_prefix=interface["comment_prefix"],
  149. )
  150. _imports = ("," + interface["line_separator"] + interface["indent"]).join(interface["imports"])
  151. _comma_maybe = "," if interface["include_trailing_comma"] else ""
  152. return (
  153. f"{interface['statement']}({_line_with_comments}{interface['line_separator']}"
  154. f"{interface['indent']}{_imports}{_comma_maybe}{interface['line_separator']})"
  155. )
  156. def _vertical_grid_common(need_trailing_char: bool, **interface: Any) -> str:
  157. if not interface["imports"]:
  158. return ""
  159. interface["statement"] += (
  160. isort.comments.add_to_line(
  161. interface["comments"],
  162. "(",
  163. removed=interface["remove_comments"],
  164. comment_prefix=interface["comment_prefix"],
  165. )
  166. + interface["line_separator"]
  167. + interface["indent"]
  168. + interface["imports"].pop(0)
  169. )
  170. while interface["imports"]:
  171. next_import = interface["imports"].pop(0)
  172. next_statement = f"{interface['statement']}, {next_import}"
  173. current_line_length = len(next_statement.split(interface["line_separator"])[-1])
  174. if interface["imports"] or interface["include_trailing_comma"]:
  175. # We need to account for a comma after this import.
  176. current_line_length += 1
  177. if not interface["imports"] and need_trailing_char:
  178. # We need to account for a closing ) we're going to add.
  179. current_line_length += 1
  180. if current_line_length > interface["line_length"]:
  181. next_statement = (
  182. f"{interface['statement']},{interface['line_separator']}"
  183. f"{interface['indent']}{next_import}"
  184. )
  185. interface["statement"] = next_statement
  186. if interface["include_trailing_comma"]:
  187. interface["statement"] += ","
  188. return str(interface["statement"])
  189. @_wrap_mode
  190. def vertical_grid(**interface: Any) -> str:
  191. return _vertical_grid_common(need_trailing_char=True, **interface) + ")"
  192. @_wrap_mode
  193. def vertical_grid_grouped(**interface: Any) -> str:
  194. return (
  195. _vertical_grid_common(need_trailing_char=False, **interface)
  196. + str(interface["line_separator"])
  197. + ")"
  198. )
  199. @_wrap_mode
  200. def vertical_grid_grouped_no_comma(**interface: Any) -> str:
  201. # This is a deprecated alias for vertical_grid_grouped above. This function
  202. # needs to exist for backwards compatibility but should never get called.
  203. raise NotImplementedError
  204. @_wrap_mode
  205. def noqa(**interface: Any) -> str:
  206. _imports = ", ".join(interface["imports"])
  207. retval = f"{interface['statement']}{_imports}"
  208. comment_str = " ".join(interface["comments"])
  209. if interface["comments"]:
  210. if (
  211. len(retval) + len(interface["comment_prefix"]) + 1 + len(comment_str)
  212. <= interface["line_length"]
  213. ):
  214. return f"{retval}{interface['comment_prefix']} {comment_str}"
  215. if "NOQA" in interface["comments"]:
  216. return f"{retval}{interface['comment_prefix']} {comment_str}"
  217. return f"{retval}{interface['comment_prefix']} NOQA {comment_str}"
  218. if len(retval) <= interface["line_length"]:
  219. return retval
  220. return f"{retval}{interface['comment_prefix']} NOQA"
  221. @_wrap_mode
  222. def vertical_hanging_indent_bracket(**interface: Any) -> str:
  223. if not interface["imports"]:
  224. return ""
  225. statement = vertical_hanging_indent(**interface)
  226. return f"{statement[:-1]}{interface['indent']})"
  227. @_wrap_mode
  228. def vertical_prefix_from_module_import(**interface: Any) -> str:
  229. if not interface["imports"]:
  230. return ""
  231. prefix_statement = interface["statement"]
  232. output_statement = prefix_statement + interface["imports"].pop(0)
  233. comments = interface["comments"]
  234. statement = output_statement
  235. statement_with_comments = ""
  236. for next_import in interface["imports"]:
  237. statement = statement + ", " + next_import
  238. statement_with_comments = isort.comments.add_to_line(
  239. comments,
  240. statement,
  241. removed=interface["remove_comments"],
  242. comment_prefix=interface["comment_prefix"],
  243. )
  244. if (
  245. len(statement_with_comments.split(interface["line_separator"])[-1]) + 1
  246. > interface["line_length"]
  247. ):
  248. statement = (
  249. isort.comments.add_to_line(
  250. comments,
  251. output_statement,
  252. removed=interface["remove_comments"],
  253. comment_prefix=interface["comment_prefix"],
  254. )
  255. + f"{interface['line_separator']}{prefix_statement}{next_import}"
  256. )
  257. comments = []
  258. output_statement = statement
  259. if comments and statement_with_comments:
  260. output_statement = statement_with_comments
  261. return str(output_statement)
  262. @_wrap_mode
  263. def hanging_indent_with_parentheses(**interface: Any) -> str:
  264. if not interface["imports"]:
  265. return ""
  266. line_length_limit = interface["line_length"] - 1
  267. interface["statement"] += "("
  268. next_import = interface["imports"].pop(0)
  269. next_statement = interface["statement"] + next_import
  270. # Check for first import
  271. if len(next_statement) > line_length_limit:
  272. next_statement = (
  273. isort.comments.add_to_line(
  274. interface["comments"],
  275. interface["statement"],
  276. removed=interface["remove_comments"],
  277. comment_prefix=interface["comment_prefix"],
  278. )
  279. + f"{interface['line_separator']}{interface['indent']}{next_import}"
  280. )
  281. interface["comments"] = []
  282. interface["statement"] = next_statement
  283. while interface["imports"]:
  284. next_import = interface["imports"].pop(0)
  285. if (
  286. interface["line_separator"] not in interface["statement"]
  287. and "#" in interface["statement"]
  288. ): # pragma: no cover # TODO: fix, this is because of test run inconsistency.
  289. line, comments = interface["statement"].split("#", 1)
  290. next_statement = (
  291. f"{line.rstrip()}, {next_import}{interface['comment_prefix']}{comments}"
  292. )
  293. else:
  294. next_statement = isort.comments.add_to_line(
  295. interface["comments"],
  296. interface["statement"] + ", " + next_import,
  297. removed=interface["remove_comments"],
  298. comment_prefix=interface["comment_prefix"],
  299. )
  300. current_line = next_statement.split(interface["line_separator"])[-1]
  301. if len(current_line) > line_length_limit:
  302. next_statement = (
  303. isort.comments.add_to_line(
  304. interface["comments"],
  305. interface["statement"] + ",",
  306. removed=interface["remove_comments"],
  307. comment_prefix=interface["comment_prefix"],
  308. )
  309. + f"{interface['line_separator']}{interface['indent']}{next_import}"
  310. )
  311. interface["comments"] = []
  312. interface["statement"] = next_statement
  313. return f"{interface['statement']}{',' if interface['include_trailing_comma'] else ''})"
  314. @_wrap_mode
  315. def backslash_grid(**interface: Any) -> str:
  316. interface["indent"] = interface["white_space"][:-1]
  317. return hanging_indent(**interface)
  318. WrapModes = enum.Enum( # type: ignore
  319. "WrapModes", {wrap_mode: index for index, wrap_mode in enumerate(_wrap_modes.keys())}
  320. )