format.py 5.3 KB

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