wrap.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146
  1. import copy
  2. import re
  3. from typing import List, Optional, Sequence
  4. from .settings import DEFAULT_CONFIG, Config
  5. from .wrap_modes import WrapModes as Modes
  6. from .wrap_modes import formatter_from_string, vertical_hanging_indent
  7. def import_statement(
  8. import_start: str,
  9. from_imports: List[str],
  10. comments: Sequence[str] = (),
  11. line_separator: str = "\n",
  12. config: Config = DEFAULT_CONFIG,
  13. multi_line_output: Optional[Modes] = None,
  14. explode: bool = False,
  15. ) -> str:
  16. """Returns a multi-line wrapped form of the provided from import statement."""
  17. if explode:
  18. formatter = vertical_hanging_indent
  19. line_length = 1
  20. include_trailing_comma = True
  21. else:
  22. formatter = formatter_from_string((multi_line_output or config.multi_line_output).name)
  23. line_length = config.wrap_length or config.line_length
  24. include_trailing_comma = config.include_trailing_comma
  25. dynamic_indent = " " * (len(import_start) + 1)
  26. indent = config.indent
  27. statement = formatter(
  28. statement=import_start,
  29. imports=copy.copy(from_imports),
  30. white_space=dynamic_indent,
  31. indent=indent,
  32. line_length=line_length,
  33. comments=comments,
  34. line_separator=line_separator,
  35. comment_prefix=config.comment_prefix,
  36. include_trailing_comma=include_trailing_comma,
  37. remove_comments=config.ignore_comments,
  38. )
  39. if config.balanced_wrapping:
  40. lines = statement.split(line_separator)
  41. line_count = len(lines)
  42. if len(lines) > 1:
  43. minimum_length = min(len(line) for line in lines[:-1])
  44. else:
  45. minimum_length = 0
  46. new_import_statement = statement
  47. while len(lines[-1]) < minimum_length and len(lines) == line_count and line_length > 10:
  48. statement = new_import_statement
  49. line_length -= 1
  50. new_import_statement = formatter(
  51. statement=import_start,
  52. imports=copy.copy(from_imports),
  53. white_space=dynamic_indent,
  54. indent=indent,
  55. line_length=line_length,
  56. comments=comments,
  57. line_separator=line_separator,
  58. comment_prefix=config.comment_prefix,
  59. include_trailing_comma=include_trailing_comma,
  60. remove_comments=config.ignore_comments,
  61. )
  62. lines = new_import_statement.split(line_separator)
  63. if statement.count(line_separator) == 0:
  64. return _wrap_line(statement, line_separator, config)
  65. return statement
  66. def line(content: str, line_separator: str, config: Config = DEFAULT_CONFIG) -> str:
  67. """Returns a line wrapped to the specified line-length, if possible."""
  68. wrap_mode = config.multi_line_output
  69. if len(content) > config.line_length and wrap_mode != Modes.NOQA: # type: ignore
  70. line_without_comment = content
  71. comment = None
  72. if "#" in content:
  73. line_without_comment, comment = content.split("#", 1)
  74. for splitter in ("import ", "cimport ", ".", "as "):
  75. exp = r"\b" + re.escape(splitter) + r"\b"
  76. if re.search(exp, line_without_comment) and not line_without_comment.strip().startswith(
  77. splitter
  78. ):
  79. line_parts = re.split(exp, line_without_comment)
  80. if comment and not (config.use_parentheses and "noqa" in comment):
  81. _comma_maybe = (
  82. ","
  83. if (
  84. config.include_trailing_comma
  85. and config.use_parentheses
  86. and not line_without_comment.rstrip().endswith(",")
  87. )
  88. else ""
  89. )
  90. line_parts[
  91. -1
  92. ] = f"{line_parts[-1].strip()}{_comma_maybe}{config.comment_prefix}{comment}"
  93. next_line = []
  94. while (len(content) + 2) > (
  95. config.wrap_length or config.line_length
  96. ) and line_parts:
  97. next_line.append(line_parts.pop())
  98. content = splitter.join(line_parts)
  99. if not content:
  100. content = next_line.pop()
  101. cont_line = _wrap_line(
  102. config.indent + splitter.join(next_line).lstrip(),
  103. line_separator,
  104. config,
  105. )
  106. if config.use_parentheses:
  107. if splitter == "as ":
  108. output = f"{content}{splitter}{cont_line.lstrip()}"
  109. else:
  110. _comma = "," if config.include_trailing_comma and not comment else ""
  111. if wrap_mode in (
  112. Modes.VERTICAL_HANGING_INDENT, # type: ignore
  113. Modes.VERTICAL_GRID_GROUPED, # type: ignore
  114. ):
  115. _separator = line_separator
  116. else:
  117. _separator = ""
  118. _comment = ""
  119. if comment and "noqa" in comment:
  120. _comment = f"{config.comment_prefix}{comment}"
  121. cont_line = cont_line.rstrip()
  122. _comma = "," if config.include_trailing_comma else ""
  123. output = (
  124. f"{content}{splitter}({_comment}"
  125. f"{line_separator}{cont_line}{_comma}{_separator})"
  126. )
  127. lines = output.split(line_separator)
  128. if config.comment_prefix in lines[-1] and lines[-1].endswith(")"):
  129. content, comment = lines[-1].split(config.comment_prefix, 1)
  130. lines[-1] = content + ")" + config.comment_prefix + comment[:-1]
  131. return line_separator.join(lines)
  132. return f"{content}{splitter}\\{line_separator}{cont_line}"
  133. elif len(content) > config.line_length and wrap_mode == Modes.NOQA and "# NOQA" not in content: # type: ignore
  134. return f"{content}{config.comment_prefix} NOQA"
  135. return content
  136. _wrap_line = line