def split_frameworks()

in utils/convert_doc_to_notebooks.py [0:0]

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

def split_frameworks(code):
    """ Split code between the two frameworks (if it has two versions) with PyTorch first."""
    if _re_pytorch.search(code) is None or _re_tensorflow.search(code) is None:
        return (code,)
    lines = code.split("\n")
    is_pytorch_first = _re_pytorch.search(lines[0]) is not None
    re_split = _re_tensorflow if is_pytorch_first else _re_pytorch
    i = 1
    while re_split.search(lines[i]) is None:
        i += 1
    j = i-1
    while len(lines[j]) == 0:
        j -= 1
    return ("\n".join(lines[:j+1]), "\n".join(lines[i:])) if is_pytorch_first else ("\n".join(lines[i:]), "\n".join(lines[:j+1]))