redbaron

gistfile1.txt

import os
import redbaron
import sys


def fixup_test(test):
    if not test.arguments:
        return

    if not test.decorators:
        return

    star_found = False
    for arg in test.arguments:
        if str(arg).startswith("*"):
            star_found = True

    decorators_to_keep = []
    decorators_to_remove = []
    for decorator in test.decorators:
        if str(decorator.value).startswith("fixtures."):
            decorators_to_remove.append(decorator)
        else:
            decorators_to_keep.append(decorator)

    if not decorators_to_remove:
        return

    if star_found:
        arguments_to_keep = test.arguments[:-1]
        arguments_to_remove = [str(d.value).replace("fixtures.", "")
                               for d in decorators_to_remove]
    else:
        arguments_to_remove = test.arguments[-len(decorators_to_remove):]
        arguments_to_keep = test.arguments[:-len(decorators_to_remove)]

    if not arguments_to_remove:
        return

    test.decorators = decorators_to_keep
    test.arguments = arguments_to_keep

    code = ("with "
            + ", ".join(["%s%s as %s" % (decorator.value, decorator.call, argument)
                         for decorator, argument in
                         zip(decorators_to_remove, arguments_to_remove)])
            + ":" + test.value.dumps())

    test.value[0].replace(code)
    del test.value[1:]

    for node in test.value[0].value.find_all("endl"):
        if node.indent:
            node.indent += "    "

def fixup(path):
    with open(path, "r+") as f:
        red = redbaron.RedBaron(f.read())
        tests = red.find_all("def")

        for test in tests:
            if not test.arguments or not test.decorators:
                continue

            try:
                fixup_test(test)
            except Exception:
                print(test)
                # test.help()
                raise

        f.seek(0)
        f.write(red.dumps())


if len(sys.argv) == 2:
    fixup(sys.argv[1])
else:
    for name in os.listdir("."):
        if not name.endswith(".py"):
            continue
        if name == "test_call_logs.py":
            continue
        print("doing %s..." % name)
        fixup(name)