update salt.utils.stringutils.human_to_bytes for broader use

changelog/61831.added

@@ -0,0 +1 @@
+Add better handling for unit abbreviations and large values to salt.utils.stringutils.human_to_bytes

doc/topics/jinja/index.rst

@@ -1676,6 +1676,33 @@ Returns:
   Example 1: snakeCaseForTheWin
   Example 2: SnakeCaseForTheWin
 
+
+.. jinja_ref:: human_to_bytes
+
+``human_to_bytes``
+------------------
+
+.. versionadded:: 3005
+
+Given a human-readable byte string (e.g. 2G, 30MB, 64KiB), return the number of bytes.
+Will return 0 if the argument has unexpected form.
+
+.. code-block:: jinja
+
+  Example 1: {{ "32GB" | human_to_bytes }}
+
+  Example 2: {{ "32GB" | human_to_bytes(handle_metric=True) }}
+
+  Example 3: {{ "32" | human_to_bytes(default_unit="GiB") }}
+
+Returns:
+
+.. code-block:: text
+
+  Example 1: 34359738368
+  Example 2: 32000000000
+  Example 3: 34359738368
+
 Networking Filters
 ------------------
 

salt/modules/virt.py

@@ -911,26 +911,12 @@ def _handle_unit(s, def_unit="m"):
     """
     Handle the unit conversion, return the value in bytes
     """
-    m = re.match(r"(?P<value>[0-9.]*)\s*(?P<unit>.*)$", str(s).strip())
-    value = m.group("value")
-    # default unit
-    unit = m.group("unit").lower() or def_unit
-    try:
-        value = int(value)
-    except ValueError:
-        try:
-            value = float(value)
-        except ValueError:
-            raise SaltInvocationError("invalid number")
-    # flag for base ten
-    dec = False
-    if re.match(r"[kmgtpezy]b$", unit):
-        dec = True
-    elif not re.match(r"(b|[kmgtpezy](ib)?)$", unit):
-        raise SaltInvocationError("invalid units")
-    p = "bkmgtpezy".index(unit[0])
-    value *= 10 ** (p * 3) if dec else 2 ** (p * 10)
-    return int(value)
+    ret = salt.utils.stringutils.human_to_bytes(
+        s, default_unit=def_unit, handle_metric=True
+    )
+    if ret == 0:
+        raise SaltInvocationError("invalid number or unit")
+    return ret
 
 
 def nesthash(value=None):

salt/utils/stringutils.py

@@ -223,31 +223,36 @@ def contains_whitespace(text):
     return any(x.isspace() for x in text)
 
 
-def human_to_bytes(size):
+@jinja_filter("human_to_bytes")
+def human_to_bytes(size, default_unit="B", handle_metric=False):
     """
-    Given a human-readable byte string (e.g. 2G, 30M),
+    Given a human-readable byte string (e.g. 2G, 30MB, 64KiB),
     return the number of bytes.  Will return 0 if the argument has
     unexpected form.
 
     .. versionadded:: 2018.3.0
+    .. versionchanged:: 3005
     """
-    sbytes = size[:-1]
-    unit = size[-1]
-    if sbytes.isdigit():
-        sbytes = int(sbytes)
-        if unit == "P":
-            sbytes *= 1125899906842624
-        elif unit == "T":
-            sbytes *= 1099511627776
-        elif unit == "G":
-            sbytes *= 1073741824
-        elif unit == "M":
-            sbytes *= 1048576
-        else:
-            sbytes = 0
-    else:
-        sbytes = 0
-    return sbytes
+    m = re.match(r"(?P<value>[0-9.]*)\s*(?P<unit>.*)$", str(size).strip())
+    value = m.group("value")
+    # default unit
+    unit = m.group("unit").lower() or default_unit.lower()
+    try:
+        value = int(value)
+    except ValueError:
+        try:
+            value = float(value)
+        except ValueError:
+            return 0
+    # flag for base ten
+    dec = False
+    if re.match(r"[kmgtpezy]b$", unit):
+        dec = True if handle_metric else False
+    elif not re.match(r"(b|[kmgtpezy](ib)?)$", unit):
+        return 0
+    p = "bkmgtpezy".index(unit[0])
+    value *= 10 ** (p * 3) if dec else 2 ** (p * 10)
+    return int(value)
 
 
 def build_whitespace_split_regex(text):

tests/pytests/unit/utils/test_stringutils.py

@@ -684,3 +684,104 @@ def test_check_include_exclude_regex():
         )
         is False
     )
+
+
+@pytest.mark.parametrize(
+    "unit",
+    [
+        "B",
+        "K",
+        "KB",
+        "KiB",
+        "M",
+        "MB",
+        "MiB",
+        "G",
+        "GB",
+        "GiB",
+        "T",
+        "TB",
+        "TiB",
+        "P",
+        "PB",
+        "PiB",
+        "E",
+        "EB",
+        "EiB",
+        "Z",
+        "ZB",
+        "ZiB",
+        "Y",
+        "YB",
+        "YiB",
+    ],
+)
+def test_human_to_bytes(unit):
+    # first multiplier is IEC/binary
+    # second multiplier is metric/decimal
+    conversion = {
+        "B": (1, 1),
+        "K": (2 ** 10, 10 ** 3),
+        "M": (2 ** 20, 10 ** 6),
+        "G": (2 ** 30, 10 ** 9),
+        "T": (2 ** 40, 10 ** 12),
+        "P": (2 ** 50, 10 ** 15),
+        "E": (2 ** 60, 10 ** 18),
+        "Z": (2 ** 70, 10 ** 21),
+        "Y": (2 ** 80, 10 ** 24),
+    }
+
+    idx = 0
+    if len(unit) == 2:
+        idx = 1
+
+    # pull out the multipliers for the units
+    multiplier = conversion[unit.upper()[0]][idx]
+    iec = conversion[unit.upper()[0]][0]
+
+    vals = [32]
+    # don't calculate a half a byte
+    if unit != "B":
+        # otherwise, test a float as well
+        vals.append(64.5)
+
+    for val in vals:
+        # calculate KB, MB, GB, etc. as 1024 instead of 1000 (legacy use)
+        assert (
+            salt.utils.stringutils.human_to_bytes("{}{}".format(val, unit)) == val * iec
+        )
+        assert (
+            salt.utils.stringutils.human_to_bytes("{} {}".format(val, unit))
+            == val * iec
+        )
+        # handle metric (KB, MB, GB, etc.) per standard
+        assert (
+            salt.utils.stringutils.human_to_bytes(
+                "{}{}".format(val, unit), handle_metric=True
+            )
+            == val * multiplier
+        )
+        assert (
+            salt.utils.stringutils.human_to_bytes(
+                "{} {}".format(val, unit), handle_metric=True
+            )
+            == val * multiplier
+        )
+
+
+def test_human_to_bytes_edge_cases():
+    # no unit - bytes
+    assert salt.utils.stringutils.human_to_bytes("32") == 32
+    # no unit - default MB
+    assert salt.utils.stringutils.human_to_bytes("32", default_unit="M") == 32 * 2 ** 20
+    # bad value
+    assert salt.utils.stringutils.human_to_bytes("32-1") == 0
+    assert salt.utils.stringutils.human_to_bytes("3.4.MB") == 0
+    assert salt.utils.stringutils.human_to_bytes("") == 0
+    assert salt.utils.stringutils.human_to_bytes("bytes") == 0
+    # bad unit
+    assert salt.utils.stringutils.human_to_bytes("32gigajammers") == 0
+    assert salt.utils.stringutils.human_to_bytes("512bytes") == 0
+    assert salt.utils.stringutils.human_to_bytes("4 Kbytes") == 0
+    assert salt.utils.stringutils.human_to_bytes("9ib") == 0
+    assert salt.utils.stringutils.human_to_bytes("2HB") == 0