format.py 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157
  1. import re
  2. import sys
  3. from datetime import datetime
  4. from difflib import unified_diff
  5. from pathlib import Path
  6. from typing import Optional, TextIO
  7. try:
  8. import colorama
  9. except ImportError:
  10. colorama_unavailable = True
  11. else:
  12. colorama_unavailable = False
  13. ADDED_LINE_PATTERN = re.compile(r"\+[^+]")
  14. REMOVED_LINE_PATTERN = re.compile(r"-[^-]")
  15. def format_simplified(import_line: str) -> str:
  16. import_line = import_line.strip()
  17. if import_line.startswith("from "):
  18. import_line = import_line.replace("from ", "")
  19. import_line = import_line.replace(" import ", ".")
  20. elif import_line.startswith("import "):
  21. import_line = import_line.replace("import ", "")
  22. return import_line
  23. def format_natural(import_line: str) -> str:
  24. import_line = import_line.strip()
  25. if not import_line.startswith("from ") and not import_line.startswith("import "):
  26. if "." not in import_line:
  27. return f"import {import_line}"
  28. parts = import_line.split(".")
  29. end = parts.pop(-1)
  30. return f"from {'.'.join(parts)} import {end}"
  31. return import_line
  32. def show_unified_diff(
  33. *,
  34. file_input: str,
  35. file_output: str,
  36. file_path: Optional[Path],
  37. output: Optional[TextIO] = None,
  38. color_output: bool = False,
  39. ) -> None:
  40. """Shows a unified_diff for the provided input and output against the provided file path.
  41. - **file_input**: A string that represents the contents of a file before changes.
  42. - **file_output**: A string that represents the contents of a file after changes.
  43. - **file_path**: A Path object that represents the file path of the file being changed.
  44. - **output**: A stream to output the diff to. If non is provided uses sys.stdout.
  45. - **color_output**: Use color in output if True.
  46. """
  47. printer = create_terminal_printer(color_output, output)
  48. file_name = "" if file_path is None else str(file_path)
  49. file_mtime = str(
  50. datetime.now() if file_path is None else datetime.fromtimestamp(file_path.stat().st_mtime)
  51. )
  52. unified_diff_lines = unified_diff(
  53. file_input.splitlines(keepends=True),
  54. file_output.splitlines(keepends=True),
  55. fromfile=file_name + ":before",
  56. tofile=file_name + ":after",
  57. fromfiledate=file_mtime,
  58. tofiledate=str(datetime.now()),
  59. )
  60. for line in unified_diff_lines:
  61. printer.diff_line(line)
  62. def ask_whether_to_apply_changes_to_file(file_path: str) -> bool:
  63. answer = None
  64. while answer not in ("yes", "y", "no", "n", "quit", "q"):
  65. answer = input(f"Apply suggested changes to '{file_path}' [y/n/q]? ") # nosec
  66. answer = answer.lower()
  67. if answer in ("no", "n"):
  68. return False
  69. if answer in ("quit", "q"):
  70. sys.exit(1)
  71. return True
  72. def remove_whitespace(content: str, line_separator: str = "\n") -> str:
  73. content = content.replace(line_separator, "").replace(" ", "").replace("\x0c", "")
  74. return content
  75. class BasicPrinter:
  76. ERROR = "ERROR"
  77. SUCCESS = "SUCCESS"
  78. def __init__(self, error: str, success: str, output: Optional[TextIO] = None):
  79. self.output = output or sys.stdout
  80. self.success_message = success
  81. self.error_message = error
  82. def success(self, message: str) -> None:
  83. print(self.success_message.format(success=self.SUCCESS, message=message), file=self.output)
  84. def error(self, message: str) -> None:
  85. print(self.error_message.format(error=self.ERROR, message=message), file=sys.stderr)
  86. def diff_line(self, line: str) -> None:
  87. self.output.write(line)
  88. class ColoramaPrinter(BasicPrinter):
  89. def __init__(self, error: str, success: str, output: Optional[TextIO]):
  90. super().__init__(error, success, output=output)
  91. # Note: this constants are instance variables instead ofs class variables
  92. # because they refer to colorama which might not be installed.
  93. self.ERROR = self.style_text("ERROR", colorama.Fore.RED)
  94. self.SUCCESS = self.style_text("SUCCESS", colorama.Fore.GREEN)
  95. self.ADDED_LINE = colorama.Fore.GREEN
  96. self.REMOVED_LINE = colorama.Fore.RED
  97. @staticmethod
  98. def style_text(text: str, style: Optional[str] = None) -> str:
  99. if style is None:
  100. return text
  101. return style + text + str(colorama.Style.RESET_ALL)
  102. def diff_line(self, line: str) -> None:
  103. style = None
  104. if re.match(ADDED_LINE_PATTERN, line):
  105. style = self.ADDED_LINE
  106. elif re.match(REMOVED_LINE_PATTERN, line):
  107. style = self.REMOVED_LINE
  108. self.output.write(self.style_text(line, style))
  109. def create_terminal_printer(
  110. color: bool, output: Optional[TextIO] = None, error: str = "", success: str = ""
  111. ) -> BasicPrinter:
  112. if color and colorama_unavailable:
  113. no_colorama_message = (
  114. "\n"
  115. "Sorry, but to use --color (color_output) the colorama python package is required.\n\n"
  116. "Reference: https://pypi.org/project/colorama/\n\n"
  117. "You can either install it separately on your system or as the colors extra "
  118. "for isort. Ex: \n\n"
  119. "$ pip install isort[colors]\n"
  120. )
  121. print(no_colorama_message, file=sys.stderr)
  122. sys.exit(1)
  123. if not colorama_unavailable:
  124. colorama.init(strip=False)
  125. return (
  126. ColoramaPrinter(error, success, output) if color else BasicPrinter(error, success, output)
  127. )