X-Git-Url: https://git.madduck.net/etc/vim.git/blobdiff_plain/229498e531b26f93c482115b1ccdb16b70bad620..f1813e31b6deed0901c8a7fb1f102b9af53de351:/src/black/output.py

diff --git a/src/black/output.py b/src/black/output.py
index c253c85..f030d0a 100644
--- a/src/black/output.py
+++ b/src/black/output.py
@@ -3,6 +3,7 @@
 The double calls are for patching purposes in tests.
 """
 
+import json
 from typing import Any, Optional
 from mypy_extensions import mypyc_attr
 import tempfile
@@ -10,6 +11,7 @@ import tempfile
 from click import echo, style
 
 
+@mypyc_attr(patchable=True)
 def _out(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
     if message is not None:
         if "bold" not in styles:
@@ -18,6 +20,7 @@ def _out(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
     echo(message, nl=nl, err=True)
 
 
+@mypyc_attr(patchable=True)
 def _err(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
     if message is not None:
         if "fg" not in styles:
@@ -26,6 +29,7 @@ def _err(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
     echo(message, nl=nl, err=True)
 
 
+@mypyc_attr(patchable=True)
 def out(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
     _out(message, nl=nl, **styles)
 
@@ -34,12 +38,29 @@ def err(message: Optional[str] = None, nl: bool = True, **styles: Any) -> None:
     _err(message, nl=nl, **styles)
 
 
+def ipynb_diff(a: str, b: str, a_name: str, b_name: str) -> str:
+    """Return a unified diff string between each cell in notebooks `a` and `b`."""
+    a_nb = json.loads(a)
+    b_nb = json.loads(b)
+    diff_lines = [
+        diff(
+            "".join(a_nb["cells"][cell_number]["source"]) + "\n",
+            "".join(b_nb["cells"][cell_number]["source"]) + "\n",
+            f"{a_name}:cell_{cell_number}",
+            f"{b_name}:cell_{cell_number}",
+        )
+        for cell_number, cell in enumerate(a_nb["cells"])
+        if cell["cell_type"] == "code"
+    ]
+    return "".join(diff_lines)
+
+
 def diff(a: str, b: str, a_name: str, b_name: str) -> str:
     """Return a unified diff string between strings `a` and `b`."""
     import difflib
 
-    a_lines = [line for line in a.splitlines(keepends=True)]
-    b_lines = [line for line in b.splitlines(keepends=True)]
+    a_lines = a.splitlines(keepends=True)
+    b_lines = b.splitlines(keepends=True)
     diff_lines = []
     for line in difflib.unified_diff(
         a_lines, b_lines, fromfile=a_name, tofile=b_name, n=5