From d5083c333f4d4783f51885bb85ce4b96954f40b3 Mon Sep 17 00:00:00 2001 From: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Date: Fri, 18 Feb 2022 14:49:53 +0100 Subject: [PATCH] style_doc handles decorators in examples (#15719) --- utils/style_doc.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/utils/style_doc.py b/utils/style_doc.py index dc021446f5..eb8398c412 100644 --- a/utils/style_doc.py +++ b/utils/style_doc.py @@ -171,9 +171,14 @@ def format_code_example(code: str, max_len: int, in_docstring: bool = False): # black may have added some new lines, we remove them code_sample = code_sample.strip() in_triple_quotes = False + in_decorator = False for line in code_sample.strip().split("\n"): if has_doctest and not is_empty_line(line): - prefix = "... " if line.startswith(" ") or line in [")", "]", "}"] or in_triple_quotes else ">>> " + prefix = ( + "... " + if line.startswith(" ") or line in [")", "]", "}"] or in_triple_quotes or in_decorator + else ">>> " + ) else: prefix = "" indent_str = "" if is_empty_line(line) else (" " * indent) @@ -181,6 +186,10 @@ def format_code_example(code: str, max_len: int, in_docstring: bool = False): if '"""' in line: in_triple_quotes = not in_triple_quotes + if line.startswith(" "): + in_decorator = False + if line.startswith("@"): + in_decorator = True formatted_lines.extend([" " * indent + line for line in output.split("\n")]) if not output.endswith("===PT-TF-SPLIT==="):