Merge pull request 'feat: add --force overwrite option' (#2) from feat/force-overwrite into main

Reviewed-on: #2

README.md

@@ -16,7 +16,7 @@ Useful for storage testing, transfer checks, and dumb-fun disk abuse.
 ### Create a file
 
 ```bash
-python make_big_file.py <output> <size> [--chunk-size SIZE] [--sparse] [--quiet]
+python make_big_file.py <output> <size> [--chunk-size SIZE] [--sparse] [--force] [--quiet]
 ```
 
 Examples:
@@ -26,6 +26,7 @@ python make_big_file.py test.bin 15GB
 python make_big_file.py dump.dat 1.5TiB --chunk-size 128MB
 python make_big_file.py tiny.bin 500MB --quiet
 python make_big_file.py sparse.img 20GB --sparse
+python make_big_file.py test.bin 15GB --force
 ```
 
 ### Read a file
@@ -66,4 +67,5 @@ Plain numbers are treated as bytes.
 ## Notes
 
 - `--sparse` is handy when you want a huge file without actually burning the disk.
+- `--force` overwrites an existing output file (without it, existing files are protected).
 - `--hash` is SHA256, because anything weaker would be cosplay.

big_file_gen.py

@@ -102,11 +102,20 @@ def _maybe_log_throughput(prefix: str, bytes_done: int, started_at: float, state
     return ProgressState(state.total_bytes, state.last_reported_percent, now, bytes_done)
 
 
-def create_file(output: str | Path, total_bytes: int, chunk_size: int, quiet: bool = False, sparse: bool = False) -> int:
+def create_file(
+    output: str | Path,
+    total_bytes: int,
+    chunk_size: int,
+    quiet: bool = False,
+    sparse: bool = False,
+    force: bool = False,
+) -> int:
     output_path = Path(output)
+    existed_before = output_path.exists()
 
-    if output_path.exists():
+    if existed_before and not force:
         print(f"Error: file already exists: {output_path}", file=sys.stderr)
+        print("Use --force to overwrite.", file=sys.stderr)
         return 1
 
     _ensure_parent_dir(output_path)
@@ -166,7 +175,7 @@ def create_file(output: str | Path, total_bytes: int, chunk_size: int, quiet: bo
     except KeyboardInterrupt:
         print("Interrupted, cleaning up partial file", file=sys.stderr)
         try:
-            if output_path.exists():
+            if output_path.exists() and not existed_before:
                 output_path.unlink()
         except OSError:
             pass
@@ -174,7 +183,7 @@ def create_file(output: str | Path, total_bytes: int, chunk_size: int, quiet: bo
     except OSError as exc:
         print(f"Error writing file: {exc}", file=sys.stderr)
         try:
-            if output_path.exists():
+            if output_path.exists() and not existed_before:
                 output_path.unlink()
         except OSError:
             pass
@@ -267,6 +276,7 @@ def build_create_parser(prog: str) -> argparse.ArgumentParser:
     parser.add_argument("size", help="Target size, for example 15GB or 1.5TiB")
     parser.add_argument("--chunk-size", default="64MB", help="Write chunk size (default: 64MB)")
     parser.add_argument("--sparse", action="store_true", help="Create a sparse file instead of writing zeros")
+    parser.add_argument("--force", "-f", action="store_true", help="Overwrite output file if it already exists")
     parser.add_argument("--quiet", "-q", action="store_true", help="Suppress progress output")
     parser.add_argument("--version", action="version", version=f"{prog} {VERSION}")
     return parser
@@ -301,7 +311,7 @@ def create_main(argv: Optional[Iterable[str]] = None) -> int:
     except ValueError as exc:
         print(f"Error: {exc}", file=sys.stderr)
         return 1
-    return create_file(args.output, total_bytes, chunk_size, args.quiet, args.sparse)
+    return create_file(args.output, total_bytes, chunk_size, args.quiet, args.sparse, args.force)
 
 
 def read_main(argv: Optional[Iterable[str]] = None) -> int: