hack-vm: allow directory input_path

1 files changed, 54 insertions, 42 deletions

diff --git a/projects/hack-vm/__main__.py b/projects/hack-vm/__main__.py
index 33df369..bacbc94 100644
--- a/projects/hack-vm/__main__.py
+++ b/projects/hack-vm/__main__.py
@@ -1,3 +1,5 @@
+import os
+from pathlib import Path
 from argparse import ArgumentParser
 from .memory import translate_memory
 from .arith_logic import translate_arith_logic
@@ -8,56 +10,66 @@ from .utils import *
 
 
 def vm_translate(input_path, verbose):
-    # program name, for static variable labels
-    input_fn = input_path.split("/")[-1]
-    prog = input_fn[:-3] if input_fn.endswith(".vm") else input_fn
-
+    # if input_path is a file, translate the file
+    # if it is a directory, translate every .vm file inside
     try:
-        input_file = open(input_path)
+        filenames = os.listdir(input_path)
+        files = [Path(input_path / f) for f in filenames]
+        vm_files = filter(lambda fn: fn.suffix == ".vm", files)
+        output_path = input_path / (input_path.name + ".asm")
+    except NotADirectoryError:
+        vm_files = [Path(input_path)]
+        output_path = input_path.with_suffix(".asm")
     except:
-        exit_on_error(f"Cannot open input file: {input_fn}", EXIT_CODE_FILE_ERROR)
+        exit_on_error(f"Cannot open input path: {input_path}", EXIT_CODE_FILE_ERROR)
+
+    # assembly translated from all .vm files will be concatenated here
+    asm = ""
 
-    # load all vm commands from file
-    vm_cmds = []
-    line = input_file.readline()
-    while line:
-        cmd = line.rstrip("\n").split("//", maxsplit=1)[0].strip()
-        if cmd:
-            vm_cmds.append(cmd)
+    for input_fn in vm_files:
+        # load all vm commands from file, e.g. program/Main.vm
+        try:
+            input_file = open(input_fn)
+        except:
+            exit_on_error(f"Cannot open input file: {input_fn}", EXIT_CODE_FILE_ERROR)
+
+        prog = input_fn.stem  # e.g. Main
+        vm_cmds = []
         line = input_file.readline()
+        while line:
+            cmd = line.rstrip("\n").split("//", maxsplit=1)[0].strip()
+            if cmd:
+                vm_cmds.append(cmd)
+            line = input_file.readline()
 
-    input_file.close()
+        input_file.close()
 
-    asm = ""
-    for cmd in vm_cmds:
-        # tokenize vm_line, hand to sub-translators based on first token
-        match cmd.split():
-            case []:
-                continue
-            case [("push" | "pop") as action, segment, index]:
-                asm += translate_memory(action, segment, index, prog, verbose)
-            case [("add" | "sub" | "neg" | "and" | "or" | "not") as command]:
-                asm += translate_arith_logic(command, verbose)
-            case [("lt" | "eq" | "gt") as command]:
-                asm += translate_compare(command, verbose)
-            case [("label" | "goto" | "if-goto") as action, label]:
-                asm += translate_branching(action, label, verbose)
-            case [("call" | "function") as action, function, n]:
-                asm += translate_function(action, function, n, prog, verbose)
-            case ["return"]:
-                asm += translate_return(verbose)
-            case _:
-                exit_on_error(
-                    f"Syntax error: invalid command: {cmd}", EXIT_CODE_SYNTAX_ERROR
-                )
+        for cmd in vm_cmds:
+            # tokenize vm_line, hand to sub-translators based on first token
+            match cmd.split():
+                case []:
+                    continue
+                case [("push" | "pop") as action, segment, index]:
+                    asm += translate_memory(action, segment, index, prog, verbose)
+                case [("add" | "sub" | "neg" | "and" | "or" | "not") as command]:
+                    asm += translate_arith_logic(command, verbose)
+                case [("lt" | "eq" | "gt") as command]:
+                    asm += translate_compare(command, verbose)
+                case [("label" | "goto" | "if-goto") as action, label]:
+                    asm += translate_branching(action, label, verbose)
+                case [("call" | "function") as action, function, n]:
+                    asm += translate_function(action, function, n, prog, verbose)
+                case ["return"]:
+                    asm += translate_return(verbose)
+                case _:
+                    exit_on_error(
+                        f"Syntax error: invalid command: {cmd}", EXIT_CODE_SYNTAX_ERROR
+                    )
 
-    output_path = (
-        input_path[:-3] + ".asm" if input_path.endswith(".vm") else input_path + ".asm"
-    )
     try:
         output_file = open(output_path, "w")
     except:
-        exit_on_error(f"Cannot open output file: {output_fn}", EXIT_CODE_FILE_ERROR)
+        exit_on_error(f"Cannot open output file: {output_path}", EXIT_CODE_FILE_ERROR)
 
     output_file.write(asm)
     output_file.close()
@@ -67,6 +79,6 @@ def vm_translate(input_path, verbose):
 if __name__ == "__main__":
     parser = ArgumentParser("hack-vm")
     parser.add_argument("-v", "--verbose", action="store_true", help="verbose mode")
-    parser.add_argument("input_path", help="input file in HACK VM")
+    parser.add_argument("input_path", help="HACK VM file or directory")
     args = parser.parse_args()
-    vm_translate(args.input_path, args.verbose)
+    vm_translate(Path(args.input_path), args.verbose)