add argparse handling for `rewritegcode`

rewritegcode.py

@@ -1,3 +1,5 @@
+import argparse
+import os.path
 import re
 import sys
 from dataclasses import dataclass
@@ -110,20 +112,36 @@ def format_info_error(info):
     return output
 
 
-def rewrite_gcode(infile, outfile):
+class ExtensionFileType(argparse.FileType):
-    with open(infile) as fp:
+    def __init__(self, ext=None, **kwargs):
-        lines = fp.read().splitlines()
+        super().__init__(**kwargs)
+        self.ext = ext
+
+    def __call__(self, string):
+        if self.ext is not None:
+            _, ext = os.path.splitext(string)
+            if not ext or ext != self.ext:
+                message = f"Expected extension {self.ext} but got {ext}"
+                raise argparse.ArgumentTypeError(message)
+
+        return super().__call__(string)
+
 
+
+def rewrite_gcode(infile, outfile):
+    lines = infile.read().splitlines()
     lines, info = rewrite_lines(lines)
 
     if not check_okay(info):
         print(format_info_error(info))
         return 1
 
-    with open(outfile, "w") as fp:
+    outfile.write(lines)
-        fp.write(lines)
-
     return 0
 
 if __name__ == "__main__":
-    sys.exit(rewrite_gcode(sys.argv[1], sys.argv[2]))
+    parser = argparse.ArgumentParser()
+    parser.add_argument("infile", type=ExtensionFileType(ext=".g", mode="r"))
+    parser.add_argument("outfile", type=ExtensionFileType(ext=".g", mode="w"))
+    args = parser.parse_args()
+    sys.exit(rewrite_gcode(args.infile, args.outfile))