wrap.py 5.9 KB

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