Mindustry jump to label

jfixer.py

"""
Usage:

python jfixer.py masm.txt

Put `label` instructions in your code:

label here
...
jump here ...
"""

import argparse
import sys


parser = argparse.ArgumentParser()
parser.add_argument(
    'file', type=argparse.FileType('r'))
parser.add_argument(
    '--output', type=argparse.FileType('w'), default=sys.stdout)


def fix_jumps(lines):
    labels = {}
    program = []
    for line in lines.splitlines():
        if line.startswith('label '):
            name = line.split(' ', 1)[1]
            labels[name] = len(program)
        else:
            program.append(line)
    for index, line in enumerate(program):
        if line.startswith('jump '):
            _, label, args = line.split(' ', 2)
            program[index] = ' '.join(['jump', str(labels[label]), args])
    if max(labels.values()) >= len(program):
        program.append('end')
    return '\n'.join(program) + '\n'


def main(in_file, out_file):
    out_file.write(fix_jumps(in_file.read()))


args = parser.parse_args()
main(args.file, args.output)