mutate.py

import os

cwd = os.path.abspath(os.path.expanduser(os.path.dirname(__file__)))


def backup(frag):
    """
    Parse a string backwards to find the first value which can
    be a variable name, parenthisized or bracked expression.

    """
    start = {")": 0, "]": 0}
    end = {"[": "]", "(": ")"}

    for i in range(len(frag) - 1, 0, -1):
        c = frag[i]
        if c in start:
            start[c] += 1
            continue
        elif c in end:
            start[end[c]] -= 1

        # if we have no open bracket or parenthisis groups
        # and we have hit whitespace exit
        if all(s == 0 for s in start.values()) and c == " ":
            break

        # or if we have exited our group like an outside parenthesis
        if any(s < 0 for s in start.values()):
            break

    # split
    split = i - sum(start.values())
    a, b = frag[:split].strip(), frag[split:].strip()

    assert b.count(")") == b.count("(")
    assert b.count("]") == b.count("[")

    return a, b


# manually replace multiple occurrances
maps = {
    "assert rec.visual.uv.ptp(axis=0).ptp() > 1e-5": "assert g.np.ptp(g.np.ptp(rec.visual.uv, axis=0)) > 1e-5",
    "assert viz.vertex_colors.ptp(axis=0).ptp() != 0": "assert g.np.ptp(g.np.ptp(viz.vertex_colors, axis=0)) != 0",
    "assert m.visual.vertex_colors.ptp(axis=0).ptp() > 0": "assert g.np.ptp(g.np.ptp(m.visual.vertex_colors, axis=0)) > 0",
    "(val[0] - val[1]).ptp(axis=1), (val[0] - val[1][::-1]).ptp(axis=1)": "np.ptp(val[0] - val[1], axis=1), np.ptp(val[0] - val[1][::-1], axis=1)",
}


def ptp(line: str, imp: str = "np") -> str:
    """
    Jankily replace `ndarray.ptp` with `numpy.ptp(...)`
    """
    split = ".ptp("
    if f"{imp}.ptp" in line or split not in line:
        # skip
        return line

    # this does not support multiple occurances
    if line.count(split) > 1:
        return maps[line]

    #
    left_initial, right = line.split(split, 1)

    # get the left side of the expression and the target
    left, target = backup(left_initial)

    # any args to ptp
    args, right = right.split(")", 1)

    args = args.strip()
    if len(args) > 0:
        args = f", {args}"

    # compose the expression
    final = f"{left} {imp}.ptp({target}{args}){right}"

    assert final.count("(") == final.count(")")
    assert final.count("[") == final.count("]")
    assert "np.ptp()" not in final

    return final


if __name__ == "__main__":
    # go through every python file in the current working directory
    for path, b, names in os.walk(cwd):
        if path == cwd:
            continue

        for name in names:
            if not name.endswith(".py"):
                continue

            path_full = os.path.join(path, name)
            with open(path_full) as f:
                text = f.read()

            lines = [L.strip() for L in str.splitlines(text) if ".ptp(" in L]

            if len(lines) == 0:
                continue

            if "import numpy as np" in text:
                imp = "np"
            else:
                imp = "g.np"

            fixed = [ptp(L, imp) for L in lines]
            for fix, ori in zip(fixed, lines):
                text = text.replace(ori, fix)
                print("\n")
                print(ori)
                print(fix)

            # write the changed file
            with open(path_full, "w") as f:
                f.write(text)