From 6b5753a41781ef7a8746aeea52c26e07f4b43c27 Mon Sep 17 00:00:00 2001
From: Tom Saunders <tom-saunders@users.noreply.github.com>
Date: Sat, 5 Sep 2020 20:15:28 +0100
Subject: [PATCH] Handle .COLOR_DIFF in the same way as .DIFF (#1673)

---
 CHANGES.md            |  2 ++
 src/black/__init__.py |  6 ++---
 tests/test_black.py   | 53 ++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/CHANGES.md b/CHANGES.md
index 7e356f1..52c8016 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -12,6 +12,8 @@
 
 - fixed a crash when PWD=/ on POSIX (#1631)
 
+- Prevent coloured diff output being interleaved with multiple files (#1673)
+
 ### 20.8b1
 
 #### _Packaging_
diff --git a/src/black/__init__.py b/src/black/__init__.py
index 64a1865..f7e7603 100644
--- a/src/black/__init__.py
+++ b/src/black/__init__.py
@@ -661,7 +661,7 @@ def reformat_one(
                 changed = Changed.YES
         else:
             cache: Cache = {}
-            if write_back != WriteBack.DIFF:
+            if write_back not in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
                 cache = read_cache(mode)
                 res_src = src.resolve()
                 if res_src in cache and cache[res_src] == get_cache_info(res_src):
@@ -735,7 +735,7 @@ async def schedule_formatting(
     :func:`format_file_in_place`.
     """
     cache: Cache = {}
-    if write_back != WriteBack.DIFF:
+    if write_back not in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
         cache = read_cache(mode)
         sources, cached = filter_cached(cache, sources)
         for src in sorted(cached):
@@ -746,7 +746,7 @@ async def schedule_formatting(
     cancelled = []
     sources_to_cache = []
     lock = None
-    if write_back == WriteBack.DIFF:
+    if write_back in (WriteBack.DIFF, WriteBack.COLOR_DIFF):
         # For diff output, we need locks to ensure we don't interleave output
         # from different processes.
         manager = Manager()
diff --git a/tests/test_black.py b/tests/test_black.py
index bc80c8f..edcf720 100644
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import multiprocessing
 import asyncio
 import logging
 from concurrent.futures import ThreadPoolExecutor
@@ -1395,9 +1396,55 @@ class BlackTestCase(unittest.TestCase):
             src = (workspace / "test.py").resolve()
             with src.open("w") as fobj:
                 fobj.write("print('hello')")
-            self.invokeBlack([str(src), "--diff"])
-            cache_file = black.get_cache_file(mode)
-            self.assertFalse(cache_file.exists())
+            with patch("black.read_cache") as read_cache, patch(
+                "black.write_cache"
+            ) as write_cache:
+                self.invokeBlack([str(src), "--diff"])
+                cache_file = black.get_cache_file(mode)
+                self.assertFalse(cache_file.exists())
+                write_cache.assert_not_called()
+                read_cache.assert_not_called()
+
+    def test_no_cache_when_writeback_color_diff(self) -> None:
+        mode = DEFAULT_MODE
+        with cache_dir() as workspace:
+            src = (workspace / "test.py").resolve()
+            with src.open("w") as fobj:
+                fobj.write("print('hello')")
+            with patch("black.read_cache") as read_cache, patch(
+                "black.write_cache"
+            ) as write_cache:
+                self.invokeBlack([str(src), "--diff", "--color"])
+                cache_file = black.get_cache_file(mode)
+                self.assertFalse(cache_file.exists())
+                write_cache.assert_not_called()
+                read_cache.assert_not_called()
+
+    @event_loop()
+    def test_output_locking_when_writeback_diff(self) -> None:
+        with cache_dir() as workspace:
+            for tag in range(0, 4):
+                src = (workspace / f"test{tag}.py").resolve()
+                with src.open("w") as fobj:
+                    fobj.write("print('hello')")
+            with patch("black.Manager", wraps=multiprocessing.Manager) as mgr:
+                self.invokeBlack(["--diff", str(workspace)], exit_code=0)
+                # this isn't quite doing what we want, but if it _isn't_
+                # called then we cannot be using the lock it provides
+                mgr.assert_called()
+
+    @event_loop()
+    def test_output_locking_when_writeback_color_diff(self) -> None:
+        with cache_dir() as workspace:
+            for tag in range(0, 4):
+                src = (workspace / f"test{tag}.py").resolve()
+                with src.open("w") as fobj:
+                    fobj.write("print('hello')")
+            with patch("black.Manager", wraps=multiprocessing.Manager) as mgr:
+                self.invokeBlack(["--diff", "--color", str(workspace)], exit_code=0)
+                # this isn't quite doing what we want, but if it _isn't_
+                # called then we cannot be using the lock it provides
+                mgr.assert_called()
 
     def test_no_cache_when_stdin(self) -> None:
         mode = DEFAULT_MODE
-- 
2.39.2