From b625df62bb2549c2b7c731ea149739fda0e10879 Mon Sep 17 00:00:00 2001
From: Phuck <pillow@mailinator.com>
Date: Sun, 6 Aug 2023 21:25:46 -0400
Subject: [PATCH] add argparse handling for `rewritegcode`

---
 rewritegcode.py | 32 +++++++++++++++++++++++++-------
 1 file changed, 25 insertions(+), 7 deletions(-)

diff --git a/rewritegcode.py b/rewritegcode.py
index 94815b6..08a40e3 100644
--- a/rewritegcode.py
+++ b/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):
-    with open(infile) as fp:
-        lines = fp.read().splitlines()
+class ExtensionFileType(argparse.FileType):
+    def __init__(self, ext=None, **kwargs):
+        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:
-        fp.write(lines)
-
+    outfile.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))