From 0ecfe6c59a2b8b175a5072eb71b6d4157667f7c6 Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Sun, 1 Jun 2025 12:15:15 +0800
Subject: [PATCH] Cast numeric string in readable_to_bytes

---
 Whatsapp_Chat_Exporter/__main__.py | 16 ++++++----------
 Whatsapp_Chat_Exporter/utility.py  |  3 +++
 tests/test_utility.py              |  3 +--
 3 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py
index 1cdb5c4..f561bcb 100644
--- a/Whatsapp_Chat_Exporter/__main__.py
+++ b/Whatsapp_Chat_Exporter/__main__.py
@@ -334,16 +334,12 @@ def validate_args(parser: ArgumentParser, args) -> None:
 
     # Size validation and conversion
     if args.size is not None:
-        stripped_size = args.size.strip()
-        if stripped_size.isnumeric():
-            args.size = int(stripped_size)
-        else:
-            try:
-                args.size = readable_to_bytes(stripped_size)
-            except ValueError:
-                parser.error(
-                    "The value for --split must be pure bytes or use a proper unit (e.g., 1048576 or 1MB)"
-                )
+        try:
+            args.size = readable_to_bytes(args.size)
+        except ValueError:
+            parser.error(
+                "The value for --split must be pure bytes or use a proper unit (e.g., 1048576 or 1MB)"
+            )
 
     # Date filter validation and processing
     if args.filter_date is not None:
diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py
index c32e536..fc97d33 100644
--- a/Whatsapp_Chat_Exporter/utility.py
+++ b/Whatsapp_Chat_Exporter/utility.py
@@ -119,6 +119,9 @@ def readable_to_bytes(size_str: str) -> int:
     }
     size_str = size_str.upper().strip()
     match = re.fullmatch(r'^(\d+(\.\d+)?)\s*([KMGTPEZY]?B)?$', size_str)
+    if size_str.isnumeric():
+        # If the string is purely numeric, assume it's in bytes
+        return int(size_str)
     if not match:
         raise ValueError("Invalid size format. Expected format like '10MB', '1024GB', or '512'.")
     unit = ''.join(filter(str.isalpha, size_str)).strip()
diff --git a/tests/test_utility.py b/tests/test_utility.py
index 6a7c326..c4d57f2 100644
--- a/tests/test_utility.py
+++ b/tests/test_utility.py
@@ -76,8 +76,7 @@ class TestReadableToBytes:
             readable_to_bytes("ABC KB")
 
     def test_missing_unit(self):
-        with pytest.raises(ValueError, match="Invalid input for size_str"):
-            readable_to_bytes("100")
+        assert readable_to_bytes("100") == 100
 
 
 class TestSanitizeExcept: