Simple assembler for CS130 creature scripts

assembler.py

"""
# A primitive assembler for creature instructions.

## Usage

```
python assembler.py src.txt out.txt
```

## Source format

A label is any non-integer string with no whitespace.

Labels are defined by putting :LABEL after a line.

Anything after ; is treated as a comment and not writen to the compiled script.
Labels *must* be defined outside of comments.
"""

import sys
import logging

if len(sys.argv) != 3:
    raise Exception("Expected an input and output argument")

with open(sys.argv[1]) as in_file, open(sys.argv[2], 'w') as out_file:
    header = next(in_file)
    program = []
    labels = {}

    # Build instruction list
    for line in in_file:
        # Ignore anything after ;
        line = line.split(';')[0]
        if line.strip() == "":
            continue
        program.append(line.split())

    # Find labels
    for op in program:
        if len(op) >= 2 and op[1][0] != ":":
            try:
                int(op[1])
            # Value is not an int so must be a label
            except ValueError:
                logging.info(f"Finding label for {op[1]} ")
                # Definition already found
                if op[1] in labels:
                    continue
                # Find label definition (:label)
                for i, line in enumerate(program):
                    for token in line:
                        if token == f":{op[1]}":
                            logging.info(f"Definition found for {op[1]}")
                            # Lines are 1 indexed because of header
                            labels[op[1]] = i + 1
                            break

                if op[1] not in labels:
                    raise Exception(f"Label {op[1]} was used but no definition could be found")

    # Replace labels
    for op in program:
        if len(op) >= 2 and op[1][0] != ":":
            try:
                int(op[1])
            except ValueError:
                op[1] = str(labels[op[1]])

        # Remove labels from final output
        for i in range(len(op)):
            if op[i][0] == ":":
                op[i] = ""
    
    # Write program
    out_file.write(header)
    for line in program:
        out_file.write(" ".join(line) + "\n")

logging.info("Assembling completed")