Added bits conversion to the ByteSize class feature #8415 (#8507)

pydantic · Jan 8, 2024 · a73e1a9 · a73e1a9
1 parent e4fa099
commit a73e1a9
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 3 deletions.
diff --git a/pydantic/types.py b/pydantic/types.py
@@ -1724,6 +1724,19 @@ def validate_brand(card_number: str) -> PaymentCardBrand:
     'tib': 2**40,
     'pib': 2**50,
     'eib': 2**60,
+    'bit': 1 / 8,
+    'kbit': 10**3 / 8,
+    'mbit': 10**6 / 8,
+    'gbit': 10**9 / 8,
+    'tbit': 10**12 / 8,
+    'pbit': 10**15 / 8,
+    'ebit': 10**18 / 8,
+    'kibit': 2**10 / 8,
+    'mibit': 2**20 / 8,
+    'gibit': 2**30 / 8,
+    'tibit': 2**40 / 8,
+    'pibit': 2**50 / 8,
+    'eibit': 2**60 / 8,
 }
 BYTE_SIZES.update({k.lower()[0]: v for k, v in BYTE_SIZES.items() if 'i' not in k})
 byte_string_re = re.compile(r'^\s*(\d*\.?\d+)\s*(\w+)?', re.IGNORECASE)
@@ -1821,11 +1834,13 @@ def human_readable(self, decimal: bool = False) -> str:
         return f'{num:0.1f}{final_unit}'
 
     def to(self, unit: str) -> float:
-        """Converts a byte size to another unit.
+        """Converts a byte size to another unit, including both byte and bit units.
 
         Args:
-            unit: The unit to convert to. Must be one of the following: B, KB, MB, GB, TB, PB, EiB,
-                KiB, MiB, GiB, TiB, PiB, EiB.
+            unit: The unit to convert to. Must be one of the following: B, KB, MB, GB, TB, PB, EB,
+                KiB, MiB, GiB, TiB, PiB, EiB (byte units) and
+                bit, kbit, mbit, gbit, tbit, pbit, ebit,
+                kibit, mibit, gibit, tibit, pibit, eibit (bit units).
 
         Returns:
             The byte size in the new unit.

diff --git a/pydantic/v1/types.py b/pydantic/v1/types.py
@@ -1082,6 +1082,19 @@ def _get_brand(card_number: str) -> PaymentCardBrand:
     'tib': 2**40,
     'pib': 2**50,
     'eib': 2**60,
+    'bit': 1/8,
+    'kbit': 10**3/8,
+    'mbit': 10**6/8,
+    'gbit': 10**9/8,
+    'tbit': 10**12/8,
+    'pbit': 10**15/8,
+    'ebit': 10**18/8,
+    'kibit': 2**10/8,
+    'mibit': 2**20/8,
+    'gibit': 2**30/8,
+    'tibit': 2**40/8,
+    'pibit': 2**50/8,
+    'eibit': 2**60/8,
 }
 BYTE_SIZES.update({k.lower()[0]: v for k, v in BYTE_SIZES.items() if 'i' not in k})
 byte_string_re = re.compile(r'^\s*(\d*\.?\d+)\s*(\w+)?', re.IGNORECASE)

diff --git a/tests/test_types.py b/tests/test_types.py
@@ -4444,6 +4444,8 @@ class FrozenSetModel(BaseModel):
         ('1.5 M', int(1.5e6), '1.4MiB', '1.5MB'),
         ('5.1kib', 5222, '5.1KiB', '5.2KB'),
         ('6.2EiB', 7148113328562451456, '6.2EiB', '7.1EB'),
+        ('8bit', 1, '1B', '1B'),
+        ('1kbit', 125, '125B', '125B'),
     ),
 )
 def test_bytesize_conversions(input_value, output, human_bin, human_dec):
@@ -4467,6 +4469,8 @@ class Model(BaseModel):
     assert m.size.to('MiB') == pytest.approx(1024)
     assert m.size.to('MB') == pytest.approx(1073.741824)
     assert m.size.to('TiB') == pytest.approx(0.0009765625)
+    assert m.size.to('bit') == pytest.approx(8589934592)
+    assert m.size.to('kbit') == pytest.approx(8589934.592)
 
 
 def test_bytesize_raises():
@@ -4487,6 +4491,9 @@ class Model(BaseModel):
     with pytest.raises(PydanticCustomError, match='byte unit'):
         m.size.to('bad_unit')
 
+    with pytest.raises(PydanticCustomError, match='byte unit'):
+        m.size.to('1ZiB')
+
 
 def test_deque_success():
     class Model(BaseModel):