wrap_modes.py 13 KB

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