Add --overwrite option to CLI

OpenCageData · Aug 17, 2024 · ed3c178 · ed3c178
1 parent ee21d78
commit ed3c178
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 17 deletions.
diff --git a/opencage/command_line.py b/opencage/command_line.py
@@ -22,35 +22,48 @@ def main(args=sys.argv[1:]):
 
 def parse_args(args):
     parser = argparse.ArgumentParser(description="Opencage CLI")
+    output_mode = parse_output_mode(args)
 
     subparsers = parser.add_subparsers(dest='command')
     subparsers.required = True
 
     subparser_forward = subparsers.add_parser('forward', help="Forward geocode a file (address to coordinates)")
     subparser_reverse = subparsers.add_parser('reverse', help="Reverse geocode a file (coordinates to address)")
 
-    for sp in [subparser_forward, subparser_reverse]:
+    for subparser in [subparser_forward, subparser_reverse]:
+        subparser.add_argument("--api-key", required=True, type=api_key_type, help="your OpenCage API key")
+        subparser.add_argument("--input", required=True, type=argparse.FileType('r', encoding='utf-8'), help="input file name (one query per line)")
+        subparser.add_argument("--output", required=True, type=argparse.FileType(output_mode, encoding='utf-8'), help="output file name")
 
-        sp.add_argument("--api-key", required=True, type=api_key_type, help="your OpenCage API key")
-        sp.add_argument("--input", required=True, type=argparse.FileType('r', encoding='utf-8'), help="input file name (one query per line)")
-        sp.add_argument("--output", required=True, type=argparse.FileType('x', encoding='utf-8'), help="output file name")
+        add_optional_arguments(subparser)
 
-        sp.add_argument("--limit", type=int, default=0, help="number of lines to read from the input file")
-        sp.add_argument("--has-headers", action="store_true", help="if the first row should be treated as a header row")
-        sp.add_argument("--input-columns", type=comma_separated_type(int), default="", help="comma separated list of integers")
-        sp.add_argument("--add-columns", type=comma_separated_type(str), default="_type,country,county,city,postcode,road,house_number,confidence,formatted", help="comma separated list of output columns")
-        sp.add_argument("--workers", type=ranged_type(int, 1, 20), default=1, help="number of parallel workers")
-        sp.add_argument("--timeout", type=ranged_type(int, 1, 60), default=1, help="timeout in seconds")
-        sp.add_argument("--retries", type=ranged_type(int, 1, 60), default=10, help="number of retries")
-        sp.add_argument("--dry-run", action="store_true", help="only parse the input file")
-        sp.add_argument("--api-domain", type=str, default="api.opencagedata.com", help="API domain")
-        sp.add_argument("--extra-params", type=comma_separated_dict_type, default="", help="extra parameters appended to request URLs")
-        sp.add_argument("--no-progress", action="store_true", help="no progress bar")
-        sp.add_argument("--quiet", action="store_true", help="skip the 'all done' message")
+    return parser.parse_args(args)
 
 
-    return parser.parse_args(args)
+def parse_output_mode(args):
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--overwrite", action="store_true")
+    (parsed, _) = parser.parse_known_args(args)
+
+    return 'w' if parsed.overwrite else 'x'
+
+
+def add_optional_arguments(parser):
+    parser.add_argument("--limit", type=int, default=0, help="number of lines to read from the input file")
+    parser.add_argument("--has-headers", action="store_true", help="if the first row should be treated as a header row")
+    parser.add_argument("--input-columns", type=comma_separated_type(int), default="", help="comma separated list of integers")
+    parser.add_argument("--add-columns", type=comma_separated_type(str), default="_type,country,county,city,postcode,road,house_number,confidence,formatted", help="comma separated list of output columns")
+    parser.add_argument("--workers", type=ranged_type(int, 1, 20), default=1, help="number of parallel workers")
+    parser.add_argument("--timeout", type=ranged_type(int, 1, 60), default=1, help="timeout in seconds")
+    parser.add_argument("--retries", type=ranged_type(int, 1, 60), default=10, help="number of retries")
+    parser.add_argument("--dry-run", action="store_true", help="only parse the input file")
+    parser.add_argument("--api-domain", type=str, default="api.opencagedata.com", help="API domain")
+    parser.add_argument("--extra-params", type=comma_separated_dict_type, default="", help="extra parameters appended to request URLs")
+    parser.add_argument("--no-progress", action="store_true", help="no progress bar")
+    parser.add_argument("--quiet", action="store_true", help="skip the 'all done' message")
+    parser.add_argument("--overwrite", action="store_true", help="delete the output file first if it exists")
 
+    return parser
 
 def api_key_type(apikey):
     pattern = re.compile(r"^(oc_gc_)?[0-9a-f]{32}$")

diff --git a/test/cli/test_cli_args.py b/test/cli/test_cli_args.py
@@ -132,3 +132,23 @@ def test_defaults():
     assert args.extra_params == {}
     assert args.no_progress is False
     assert args.quiet is False
+
+def test_overwrite():
+    args = parse_args([
+        "forward",
+        "--api-key", "12345678901234567890123456789012",
+        "--input", "test/fixtures/input.txt",
+        "--output", "test/fixtures/output.csv"
+    ])
+
+    assert args.output.mode == "x"
+
+    args = parse_args([
+        "forward",
+        "--api-key", "12345678901234567890123456789012",
+        "--input", "test/fixtures/input.txt",
+        "--output", "test/fixtures/output.csv",
+        "--overwrite"
+    ])
+
+    assert args.output.mode == "w"