def format_code_example()

in src/doc_builder/style_doc.py [0:0]

• 60 lines of code
• 13 McCabe index (conditional complexity)

def format_code_example(code: str, max_len: int, in_docstring: bool = False):
    """
    Format a code example using black. Will take into account the doctest syntax as well as any initial indentation in
    the code provided.

    Args:
        code (`str`): The code example to format.
        max_len (`int`): The maximum length per line.
        in_docstring (`bool`, *optional*, defaults to `False`): Whether or not the code example is inside a docstring.

    Returns:
        `str`: The formatted code.
    """
    code_lines = code.split("\n")

    # Find initial indent
    idx = 0
    while idx < len(code_lines) and is_empty_line(code_lines[idx]):
        idx += 1
    if idx >= len(code_lines):
        return "", ""
    indent = find_indent(code_lines[idx])

    # Remove the initial indent for now, we will had it back after styling.
    # Note that l[indent:] works for empty lines
    code_lines = [l[indent:] for l in code_lines[idx:]]
    has_doctest = code_lines[0][:3] in DOCTEST_PROMPTS

    code_samples, outputs = parse_code_example(code_lines)

    # Let's blackify the code! We put everything in one big text to go faster.
    delimiter = "\n\n### New code sample ###\n"
    full_code = delimiter.join(code_samples)
    line_length = max_len - indent
    if has_doctest:
        line_length -= 4

    black_avoid_patterns = get_black_avoid_patterns()
    for k, v in black_avoid_patterns.items():
        full_code = full_code.replace(k, v)
    try:
        mode = black.Mode(target_versions={black.TargetVersion.PY37}, line_length=line_length)
        formatted_code = black.format_str(full_code, mode=mode)
        error = ""
    except Exception as e:
        formatted_code = full_code
        error = f"Code sample:\n{full_code}\n\nError message:\n{e}"

    # Let's get back the formatted code samples
    for k, v in black_avoid_patterns.items():
        formatted_code = formatted_code.replace(v, k)
    # Triple quotes will mess docstrings.
    if in_docstring:
        formatted_code = formatted_code.replace('"""', "'''")

    code_samples = formatted_code.split(delimiter)
    # We can have one output less than code samples
    if len(outputs) == len(code_samples) - 1:
        outputs.append("")

    formatted_lines = []
    for code_sample, output in zip(code_samples, outputs):
        # 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[0] in [")", "]", "}"] or in_triple_quotes or in_decorator
                    else ">>> "
                )
            else:
                prefix = ""
            indent_str = "" if is_empty_line(line) else (" " * indent)
            formatted_lines.append(indent_str + prefix + line)

            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==="):
            formatted_lines.append("")

    result = "\n".join(formatted_lines)
    return result.rstrip(), error