style_doc handles decorators in examples (#15719)

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==="):