]> git.madduck.net Git - etc/vim.git/blobdiff - black.py

madduck's git repository

Every one of the projects in this repository is available at the canonical URL git://git.madduck.net/madduck/pub/<projectpath> — see each project's metadata for the exact URL.

All patches and comments are welcome. Please squash your changes to logical commits before using git-format-patch and git-send-email to patches@git.madduck.net. If you'd read over the Git project's submission guidelines and adhered to them, I'd be especially grateful.

SSH access, as well as push access can be individually arranged.

If you use my repositories frequently, consider adding the following snippet to ~/.gitconfig and using the third clone URL listed for each project:

[url "git://git.madduck.net/madduck/"]
  insteadOf = madduck:

Add Black compatible configurations in documentation (#1366 & #1205) (#1371)
[etc/vim.git] / black.py
index 087683a34f5b149ea351979d800faf07917ee43c..26687472c597a437b487715dd7b97da76d84550a 100644 (file)
--- a/black.py
+++ b/black.py
@@ -2,7 +2,7 @@ import ast
 import asyncio
 from abc import ABC, abstractmethod
 from collections import defaultdict
 import asyncio
 from abc import ABC, abstractmethod
 from collections import defaultdict
-from concurrent.futures import Executor, ProcessPoolExecutor
+from concurrent.futures import Executor, ThreadPoolExecutor, ProcessPoolExecutor
 from contextlib import contextmanager
 from datetime import datetime
 from enum import Enum
 from contextlib import contextmanager
 from datetime import datetime
 from enum import Enum
@@ -39,6 +39,7 @@ from typing import (
     TypeVar,
     Union,
     cast,
     TypeVar,
     Union,
     cast,
+    TYPE_CHECKING,
 )
 from typing_extensions import Final
 from mypy_extensions import mypyc_attr
 )
 from typing_extensions import Final
 from mypy_extensions import mypyc_attr
@@ -59,6 +60,9 @@ from blib2to3.pgen2.parse import ParseError
 
 from _black_version import version as __version__
 
 
 from _black_version import version as __version__
 
+if TYPE_CHECKING:
+    import colorama  # noqa: F401
+
 DEFAULT_LINE_LENGTH = 88
 DEFAULT_EXCLUDES = r"/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist)/"  # noqa: B950
 DEFAULT_INCLUDES = r"\.pyi?$"
 DEFAULT_LINE_LENGTH = 88
 DEFAULT_EXCLUDES = r"/(\.eggs|\.git|\.hg|\.mypy_cache|\.nox|\.tox|\.venv|\.svn|_build|buck-out|build|dist)/"  # noqa: B950
 DEFAULT_INCLUDES = r"\.pyi?$"
@@ -140,12 +144,18 @@ class WriteBack(Enum):
     YES = 1
     DIFF = 2
     CHECK = 3
     YES = 1
     DIFF = 2
     CHECK = 3
+    COLOR_DIFF = 4
 
     @classmethod
 
     @classmethod
-    def from_configuration(cls, *, check: bool, diff: bool) -> "WriteBack":
+    def from_configuration(
+        cls, *, check: bool, diff: bool, color: bool = False
+    ) -> "WriteBack":
         if check and not diff:
             return cls.CHECK
 
         if check and not diff:
             return cls.CHECK
 
+        if diff and color:
+            return cls.COLOR_DIFF
+
         return cls.DIFF if diff else cls.YES
 
 
         return cls.DIFF if diff else cls.YES
 
 
@@ -380,6 +390,11 @@ def target_version_option_callback(
     is_flag=True,
     help="Don't write the files back, just output a diff for each file on stdout.",
 )
     is_flag=True,
     help="Don't write the files back, just output a diff for each file on stdout.",
 )
+@click.option(
+    "--color/--no-color",
+    is_flag=True,
+    help="Show colored diff. Only applies when `--diff` is given.",
+)
 @click.option(
     "--fast/--safe",
     is_flag=True,
 @click.option(
     "--fast/--safe",
     is_flag=True,
@@ -458,6 +473,7 @@ def main(
     target_version: List[TargetVersion],
     check: bool,
     diff: bool,
     target_version: List[TargetVersion],
     check: bool,
     diff: bool,
+    color: bool,
     fast: bool,
     pyi: bool,
     py36: bool,
     fast: bool,
     pyi: bool,
     py36: bool,
@@ -470,7 +486,7 @@ def main(
     config: Optional[str],
 ) -> None:
     """The uncompromising code formatter."""
     config: Optional[str],
 ) -> None:
     """The uncompromising code formatter."""
-    write_back = WriteBack.from_configuration(check=check, diff=diff)
+    write_back = WriteBack.from_configuration(check=check, diff=diff, color=color)
     if target_version:
         if py36:
             err("Cannot use both --target-version and --py36")
     if target_version:
         if py36:
             err("Cannot use both --target-version and --py36")
@@ -597,12 +613,21 @@ def reformat_many(
     sources: Set[Path], fast: bool, write_back: WriteBack, mode: Mode, report: "Report"
 ) -> None:
     """Reformat multiple files using a ProcessPoolExecutor."""
     sources: Set[Path], fast: bool, write_back: WriteBack, mode: Mode, report: "Report"
 ) -> None:
     """Reformat multiple files using a ProcessPoolExecutor."""
+    executor: Executor
     loop = asyncio.get_event_loop()
     worker_count = os.cpu_count()
     if sys.platform == "win32":
         # Work around https://bugs.python.org/issue26903
         worker_count = min(worker_count, 61)
     loop = asyncio.get_event_loop()
     worker_count = os.cpu_count()
     if sys.platform == "win32":
         # Work around https://bugs.python.org/issue26903
         worker_count = min(worker_count, 61)
-    executor = ProcessPoolExecutor(max_workers=worker_count)
+    try:
+        executor = ProcessPoolExecutor(max_workers=worker_count)
+    except OSError:
+        # we arrive here if the underlying system does not support multi-processing
+        # like in AWS Lambda, in which case we gracefully fallback to
+        # a ThreadPollExecutor with just a single worker (more workers would not do us
+        # any good due to the Global Interpreter Lock)
+        executor = ThreadPoolExecutor(max_workers=1)
+
     try:
         loop.run_until_complete(
             schedule_formatting(
     try:
         loop.run_until_complete(
             schedule_formatting(
@@ -617,7 +642,8 @@ def reformat_many(
         )
     finally:
         shutdown(loop)
         )
     finally:
         shutdown(loop)
-        executor.shutdown()
+        if executor is not None:
+            executor.shutdown()
 
 
 async def schedule_formatting(
 
 
 async def schedule_formatting(
@@ -718,12 +744,15 @@ def format_file_in_place(
     if write_back == WriteBack.YES:
         with open(src, "w", encoding=encoding, newline=newline) as f:
             f.write(dst_contents)
     if write_back == WriteBack.YES:
         with open(src, "w", encoding=encoding, newline=newline) as f:
             f.write(dst_contents)
-    elif write_back == WriteBack.DIFF:
+    elif write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
         now = datetime.utcnow()
         src_name = f"{src}\t{then} +0000"
         dst_name = f"{src}\t{now} +0000"
         diff_contents = diff(src_contents, dst_contents, src_name, dst_name)
 
         now = datetime.utcnow()
         src_name = f"{src}\t{then} +0000"
         dst_name = f"{src}\t{now} +0000"
         diff_contents = diff(src_contents, dst_contents, src_name, dst_name)
 
+        if write_back == write_back.COLOR_DIFF:
+            diff_contents = color_diff(diff_contents)
+
         with lock or nullcontext():
             f = io.TextIOWrapper(
                 sys.stdout.buffer,
         with lock or nullcontext():
             f = io.TextIOWrapper(
                 sys.stdout.buffer,
@@ -731,12 +760,57 @@ def format_file_in_place(
                 newline=newline,
                 write_through=True,
             )
                 newline=newline,
                 write_through=True,
             )
+            f = wrap_stream_for_windows(f)
             f.write(diff_contents)
             f.detach()
 
     return True
 
 
             f.write(diff_contents)
             f.detach()
 
     return True
 
 
+def color_diff(contents: str) -> str:
+    """Inject the ANSI color codes to the diff."""
+    lines = contents.split("\n")
+    for i, line in enumerate(lines):
+        if line.startswith("+++") or line.startswith("---"):
+            line = "\033[1;37m" + line + "\033[0m"  # bold white, reset
+        if line.startswith("@@"):
+            line = "\033[36m" + line + "\033[0m"  # cyan, reset
+        if line.startswith("+"):
+            line = "\033[32m" + line + "\033[0m"  # green, reset
+        elif line.startswith("-"):
+            line = "\033[31m" + line + "\033[0m"  # red, reset
+        lines[i] = line
+    return "\n".join(lines)
+
+
+def wrap_stream_for_windows(
+    f: io.TextIOWrapper,
+) -> Union[io.TextIOWrapper, "colorama.AnsiToWin32.AnsiToWin32"]:
+    """
+    Wrap the stream in colorama's wrap_stream so colors are shown on Windows.
+
+    If `colorama` is not found, then no change is made. If `colorama` does
+    exist, then it handles the logic to determine whether or not to change
+    things.
+    """
+    try:
+        from colorama import initialise
+
+        # We set `strip=False` so that we can don't have to modify
+        # test_express_diff_with_color.
+        f = initialise.wrap_stream(
+            f, convert=None, strip=False, autoreset=False, wrap=True
+        )
+
+        # wrap_stream returns a `colorama.AnsiToWin32.AnsiToWin32` object
+        # which does not have a `detach()` method. So we fake one.
+        f.detach = lambda *args, **kwargs: None  # type: ignore
+    except ImportError:
+        pass
+
+    return f
+
+
 def format_stdin_to_stdout(
     fast: bool, *, write_back: WriteBack = WriteBack.NO, mode: Mode
 ) -> bool:
 def format_stdin_to_stdout(
     fast: bool, *, write_back: WriteBack = WriteBack.NO, mode: Mode
 ) -> bool:
@@ -762,11 +836,15 @@ def format_stdin_to_stdout(
         )
         if write_back == WriteBack.YES:
             f.write(dst)
         )
         if write_back == WriteBack.YES:
             f.write(dst)
-        elif write_back == WriteBack.DIFF:
+        elif write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
             now = datetime.utcnow()
             src_name = f"STDIN\t{then} +0000"
             dst_name = f"STDOUT\t{now} +0000"
             now = datetime.utcnow()
             src_name = f"STDIN\t{then} +0000"
             dst_name = f"STDOUT\t{now} +0000"
-            f.write(diff(src, dst, src_name, dst_name))
+            d = diff(src, dst, src_name, dst_name)
+            if write_back == WriteBack.COLOR_DIFF:
+                d = color_diff(d)
+                f = wrap_stream_for_windows(f)
+            f.write(d)
         f.detach()
 
 
         f.detach()
 
 
@@ -5156,7 +5234,7 @@ def generate_ignored_nodes(leaf: Leaf) -> Iterator[LN]:
     """
     container: Optional[LN] = container_of(leaf)
     while container is not None and container.type != token.ENDMARKER:
     """
     container: Optional[LN] = container_of(leaf)
     while container is not None and container.type != token.ENDMARKER:
-        if fmt_on(container):
+        if is_fmt_on(container):
             return
 
         # fix for fmt: on in children
             return
 
         # fix for fmt: on in children
@@ -5170,17 +5248,21 @@ def generate_ignored_nodes(leaf: Leaf) -> Iterator[LN]:
             container = container.next_sibling
 
 
             container = container.next_sibling
 
 
-def fmt_on(container: LN) -> bool:
-    is_fmt_on = False
+def is_fmt_on(container: LN) -> bool:
+    """Determine whether formatting is switched on within a container.
+    Determined by whether the last `# fmt:` comment is `on` or `off`.
+    """
+    fmt_on = False
     for comment in list_comments(container.prefix, is_endmarker=False):
         if comment.value in FMT_ON:
     for comment in list_comments(container.prefix, is_endmarker=False):
         if comment.value in FMT_ON:
-            is_fmt_on = True
+            fmt_on = True
         elif comment.value in FMT_OFF:
         elif comment.value in FMT_OFF:
-            is_fmt_on = False
-    return is_fmt_on
+            fmt_on = False
+    return fmt_on
 
 
 def contains_fmt_on_at_column(container: LN, column: int) -> bool:
 
 
 def contains_fmt_on_at_column(container: LN, column: int) -> bool:
+    """Determine if children at a given column have formatting switched on."""
     for child in container.children:
         if (
             isinstance(child, Node)
     for child in container.children:
         if (
             isinstance(child, Node)
@@ -5188,13 +5270,14 @@ def contains_fmt_on_at_column(container: LN, column: int) -> bool:
             or isinstance(child, Leaf)
             and child.column == column
         ):
             or isinstance(child, Leaf)
             and child.column == column
         ):
-            if fmt_on(child):
+            if is_fmt_on(child):
                 return True
 
     return False
 
 
 def first_leaf_column(node: Node) -> Optional[int]:
                 return True
 
     return False
 
 
 def first_leaf_column(node: Node) -> Optional[int]:
+    """Returns the column of the first leaf child of a node."""
     for child in node.children:
         if isinstance(child, Leaf):
             return child.column
     for child in node.children:
         if isinstance(child, Leaf):
             return child.column