From b2d64230199ebe62fcc84e12c2dc2e3b5d485e69 Mon Sep 17 00:00:00 2001
From: Scott Gasch <scott@gasch.org>
Date: Fri, 3 Mar 2023 16:01:40 -0800
Subject: [PATCH] Add and use valid_byte_count.

---
 src/pyutils/argparse_utils.py | 65 +++++++++++++++++++++++++++++++++++
 src/pyutils/logging_utils.py  |  2 +-
 src/pyutils/string_utils.py   |  5 +++
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/pyutils/argparse_utils.py b/src/pyutils/argparse_utils.py
index 6553e2c..06b9a9c 100644
--- a/src/pyutils/argparse_utils.py
+++ b/src/pyutils/argparse_utils.py
@@ -459,6 +459,71 @@ def valid_duration(txt: str) -> datetime.timedelta:
         raise argparse.ArgumentTypeError(e) from e
 
 
+def valid_byte_count(txt: str) -> int:
+    """If the string is a valid number of bytes, return an integer
+    representing the requested byte count.  This method uses
+    :meth:`string_utils.suffix_string_to_number` to parse and and
+    accepts / understands:
+
+        - plain numbers (123456)
+        - numbers with ISO suffixes (Mb, Gb, Pb, etc...)
+
+    If the byte count is not parsable, raise an ArgumentTypeError.
+
+    Args:
+        txt: data passed to a commandline arg expecting a duration.
+
+    Returns:
+        An integer number of bytes or raises ArgumentTypeError on
+        error.
+
+    Sample usage::
+
+        cfg.add_argument(
+            '--max_file_size',
+            type=argparse_utils.valid_byte_count,
+            default=(1024 * 1024),
+            metavar='NUM_BYTES',
+            help='The largest file we may create',
+        )
+
+    >>> valid_byte_count('1Mb')
+    1048576
+
+    >>> valid_byte_count("1234567")
+    1234567
+
+    >>> valid_byte_count("1.2Gb")
+    1288490188
+
+    >>> valid_byte_count('1.2')      # <--- contains a decimal
+    Traceback (most recent call last):
+    ...
+    argparse.ArgumentTypeError: Invalid byte count: 1.2
+
+    >>> valid_byte_count(1234567)    # <--- not a string
+    Traceback (most recent call last):
+    ...
+    argparse.ArgumentTypeError: Invalid byte count: 1234567
+
+    >>> valid_byte_count('On a dark and stormy night')
+    Traceback (most recent call last):
+    ...
+    argparse.ArgumentTypeError: Invalid byte count: On a dark and stormy night
+
+    """
+    from pyutils.string_utils import suffix_string_to_number
+
+    try:
+        num_bytes = suffix_string_to_number(txt)
+        if num_bytes:
+            return num_bytes
+        raise argparse.ArgumentTypeError(f"Invalid byte count: {txt}")
+    except Exception as e:
+        logger.exception("Exception while parsing a supposed byte count: %s", txt)
+        raise argparse.ArgumentTypeError(e) from e
+
+
 if __name__ == "__main__":
     import doctest
 
diff --git a/src/pyutils/logging_utils.py b/src/pyutils/logging_utils.py
index 669ba21..8011604 100644
--- a/src/pyutils/logging_utils.py
+++ b/src/pyutils/logging_utils.py
@@ -107,7 +107,7 @@ cfg.add_argument(
 )
 cfg.add_argument(
     "--logging_filename_maxsize",
-    type=int,
+    type=argparse_utils.valid_byte_count,
     default=(1024 * 1024),
     metavar="#BYTES",
     help="The maximum size (in bytes) to write to the logging_filename.",
diff --git a/src/pyutils/string_utils.py b/src/pyutils/string_utils.py
index a503472..0c18dcc 100644
--- a/src/pyutils/string_utils.py
+++ b/src/pyutils/string_utils.py
@@ -827,6 +827,11 @@ def suffix_string_to_number(in_str: str) -> Optional[int]:
     1048576
     >>> suffix_string_to_number('13.1Gb')
     14066017894
+    >>> suffix_string_to_number('12345')
+    12345
+    >>> x = suffix_string_to_number('a lot')
+    >>> x is None
+    True
     """
 
     def suffix_capitalize(s: str) -> str:
-- 
2.45.0