diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py index 4f6f658..e601ab1 100644 --- a/Whatsapp_Chat_Exporter/__main__.py +++ b/Whatsapp_Chat_Exporter/__main__.py @@ -17,7 +17,7 @@ else: from Whatsapp_Chat_Exporter import exported_handler, android_handler from Whatsapp_Chat_Exporter import ios_handler, ios_media_handler from Whatsapp_Chat_Exporter.data_model import ChatStore -from Whatsapp_Chat_Exporter.utility import APPLE_TIME, Crypt, DbType, chat_is_empty +from Whatsapp_Chat_Exporter.utility import APPLE_TIME, Crypt, DbType, chat_is_empty, convert_size_reverse from Whatsapp_Chat_Exporter.utility import check_update, import_from_json, sanitize_filename from argparse import ArgumentParser, SUPPRESS from datetime import datetime @@ -162,7 +162,6 @@ def main(): "--split", dest="size", nargs='?', - type=int, const=0, default=None, help="Maximum (rough) size of a single output file in bytes, 0 for auto" @@ -328,6 +327,11 @@ def main(): parser.error("When --per-chat is enabled, the destination of --json must be a directory.") if args.enrich_from_vcards is not None and args.default_contry_code is None: parser.error("When --enrich-from-vcards is provided, you must also set --default-country-code") + if args.size is not None and not isinstance(args.size, int) and not args.size.isnumeric(): + try: + args.size = convert_size_reverse(args.size) + except ValueError: + parser.error("The value for --split must be ended in pure bytes or with a proper unit (e.g., 1048576 or 1MB)") if args.filter_date is not None: if " - " in args.filter_date: start, end = args.filter_date.split(" - ") diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py index af636b3..2a188ec 100644 --- a/Whatsapp_Chat_Exporter/utility.py +++ b/Whatsapp_Chat_Exporter/utility.py @@ -59,6 +59,25 @@ def convert_size(size_bytes): return "%s %s" % (s, size_name[i]) +def convert_size_reverse(size_str): + SIZE_UNITS = { + 'B': 1, + 'KB': 1024, + 'MB': 1024**2, + 'GB': 1024**3, + 'TB': 1024**4, + 'PB': 1024**5, + 'EB': 1024**6, + 'ZB': 1024**7, + 'YB': 1024**8 + } + size_str = size_str.upper().strip() + number, unit = size_str[:-2].strip(), size_str[-2:].strip() + if unit not in SIZE_UNITS or not number.isnumeric(): + raise ValueError("Invalid input for size_str. Example: 1024GB") + return int(number) * SIZE_UNITS[unit] + + def sanitize_except(html): return Markup(sanitize(html, tags=["br"]))