From 17f3a7b6dc740495b22ffa5ce6cba839b432d9ac Mon Sep 17 00:00:00 2001 From: Seth Underwood Date: Sun, 21 Jun 2026 21:03:52 -0400 Subject: [PATCH] refactor: replace if/elif unit dispatch with _UNIT_POWERS lookup Co-Authored-By: Claude Opus 4.6 --- src/fsize/__init__.py | 16 +++------------- tests/test_fsize_init.py | 17 +++++++++++++++++ 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/src/fsize/__init__.py b/src/fsize/__init__.py index f7c90c1..d5d095b 100644 --- a/src/fsize/__init__.py +++ b/src/fsize/__init__.py @@ -194,19 +194,9 @@ def __format__(self, format_spec: str) -> str: ) # Convert the number to the appropriate unit - n = self.real - if unit == "K": - n = self.to_k() - elif unit == "M": - n = self.to_m() - elif unit == "G": - n = self.to_g() - elif unit == "T": - n = self.to_t() - elif unit == "P": - n = self.to_p() - elif unit == "E": - n = self.to_e() + if unit not in _UNIT_POWERS: + raise AssertionError(f"unhandled unit: {unit!r}") + n = self.real / self._convert ** _UNIT_POWERS[unit] log_digits = math.ceil(math.log10(n)) if n > 0 else 0 out_format_spec = ( diff --git a/tests/test_fsize_init.py b/tests/test_fsize_init.py index 0ecc5a7..42f7a54 100755 --- a/tests/test_fsize_init.py +++ b/tests/test_fsize_init.py @@ -431,3 +431,20 @@ def test_format_empty_spec(): assert format(y, "") == str(y) z = FSize(1, "KB") assert format(z, "") == str(z) + + +def test_format_all_units(): + """Test that every unit in _UNIT_POWERS works via format.""" + # Binary FSize (1 EiB) + val = FSize(1, "EiB") + assert format(val, "E") == "1" + assert format(val, "P") == "1024" + assert format(val, "T") == str(1024**2) + assert format(val, "G") == str(1024**3) + assert format(val, "M") == str(1024**4) + assert format(val, "K") == str(1024**5) + # Decimal FSize (1 EB) + val_dec = FSize(1, "EB") + assert format(val_dec, "E") == "1" + assert float(format(val_dec, "P")) == 1000 + assert float(format(val_dec, "K")) == 1000**5