From 1144d336b689d1710534b245697e41be7a168075 Mon Sep 17 00:00:00 2001
From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>
Date: Tue, 18 Jan 2022 09:16:55 -0500
Subject: [PATCH] Copies and docstring styling (#15202)

* Style docstrings when making/checking copies

* Polish
---
 utils/check_copies.py |  2 ++
 utils/style_doc.py    | 39 ++++++++++++++++++++++++++++-----------
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/utils/check_copies.py b/utils/check_copies.py
index 0512e8808b..ac9ab15973 100644
--- a/utils/check_copies.py
+++ b/utils/check_copies.py
@@ -19,6 +19,7 @@ import os
 import re
 
 import black
+from style_doc import style_docstrings_in_code
 
 
 # All paths are set with the intent you should run this script from the root of the repo with the command
@@ -130,6 +131,7 @@ def blackify(code):
     if has_indent:
         code = f"class Bla:\n{code}"
     result = black.format_str(code, mode=black.FileMode([black.TargetVersion.PY35], line_length=119))
+    result, _ = style_docstrings_in_code(result)
     return result[len("class Bla:\n") :] if has_indent else result
 
 
diff --git a/utils/style_doc.py b/utils/style_doc.py
index de15a859c7..1e8aecddc6 100644
--- a/utils/style_doc.py
+++ b/utils/style_doc.py
@@ -367,9 +367,34 @@ def style_docstring(docstring, max_len):
     return "\n".join(new_lines), "\n\n".join(black_errors)
 
 
+def style_docstrings_in_code(code, max_len=119):
+    """
+    Style all docstrings in some code.
+
+    Args:
+        code (`str`): The code in which we want to style the docstrings.
+        max_len (`int`): The maximum number of characters per line.
+
+    Returns:
+        `Tuple[str, str]`: A tuple with the clean code and the black errors (if any)
+    """
+    # fmt: off
+    splits = code.split('\"\"\"')
+    splits = [
+        (s if i % 2 == 0 or _re_doc_ignore.search(splits[i - 1]) is not None else style_docstring(s, max_len=max_len))
+        for i, s in enumerate(splits)
+    ]
+    black_errors = "\n\n".join([s[1] for s in splits if isinstance(s, tuple) and len(s[1]) > 0])
+    splits = [s[0] if isinstance(s, tuple) else s for s in splits]
+    clean_code = '\"\"\"'.join(splits)
+    # fmt: on
+
+    return clean_code, black_errors
+
+
 def style_file_docstrings(code_file, max_len=119, check_only=False):
     """
-    Style all docstrings in  a given file.
+    Style all docstrings in a given file.
 
     Args:
         code_file (`str` or `os.PathLike`): The file in which we want to style the docstring.
@@ -382,16 +407,8 @@ def style_file_docstrings(code_file, max_len=119, check_only=False):
     """
     with open(code_file, "r", encoding="utf-8", newline="\n") as f:
         code = f.read()
-    # fmt: off
-    splits = code.split('\"\"\"')
-    splits = [
-        (s if i % 2 == 0 or _re_doc_ignore.search(splits[i - 1]) is not None else style_docstring(s, max_len=max_len))
-        for i, s in enumerate(splits)
-    ]
-    black_errors = "\n\n".join([s[1] for s in splits if isinstance(s, tuple) and len(s[1]) > 0])
-    splits = [s[0] if isinstance(s, tuple) else s for s in splits]
-    clean_code = '\"\"\"'.join(splits)
-    # fmt: on
+
+    clean_code, black_errors = style_docstrings_in_code(code, max_len=max_len)
 
     diff = clean_code != code
     if not check_only and diff: