def split_frameworks()

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

• 27 lines of code
• 11 McCabe index (conditional complexity)

def split_frameworks(content):
    """
    Split a given doc content in three to extract the Mixed, PyTorch and TensorFlow content.
    """
    new_lines = {"mixed": [], "pt": [], "tf": [], "jax": []}

    content = clean_doctest_syntax(content)
    lines = content.split("\n")
    idx = 0
    while idx < len(lines):
        if lines[idx].strip() == "<frameworkcontent>":
            idx += 1
            current_lines = []
            current_framework = None
            while idx < len(lines) and lines[idx].strip() != "</frameworkcontent>":
                if _re_framework.search(lines[idx]) is not None:
                    current_framework = _re_framework.search(lines[idx]).groups()[0]
                elif current_framework is not None and lines[idx].strip() == f"</{current_framework}>":
                    new_lines[current_framework].extend(current_lines)
                    new_lines["mixed"].extend(current_lines)
                    current_framework = None
                    current_lines = []
                elif current_framework is not None:
                    current_lines.append(lines[idx])
                idx += 1
            idx += 1
        else:
            for key in new_lines.keys():
                new_lines[key].append(lines[idx])
            idx += 1
    return ["\n".join(l) for l in new_lines.values()]