__len__

inbox/uid_set.py

@@ -7,42 +7,65 @@
 MAX_UINT32 = 2**32 - 1
 
 
-def compress_ranges(iterable: "Iterable[int]") -> "Iterable[int | tuple[int, int]]":
-    sorted_iterable = sorted(iterable)
+class CompressRanges:
+    def __init__(self, iterable: "Iterable[int]"):
+        self.iterable = iterable
+        self.__length: "int | None" = None
 
-    if not sorted_iterable:
-        return
+    def __iter__(self) -> "Iterator[int | tuple[int, int]]":
+        sorted_iterable = sorted(self.iterable)
 
-    cache: "int | tuple[int, int] | None" = None
+        if not sorted_iterable:
+            self.__length = 0
+            return
 
-    for element in sorted_iterable:
-        assert 0 < element <= MAX_UINT32
+        length = 0
+        cache: "int | tuple[int, int] | None" = None
 
-        if cache is None:
-            cache = element
-            continue
+        for element in sorted_iterable:
+            assert 0 < element <= MAX_UINT32
 
-        if isinstance(cache, int):
-            if element == cache:
+            if cache is None:
+                cache = element
+                length += 1
                 continue
 
-            if element - cache == 1:
-                cache = (cache, element)
-                continue
+            if isinstance(cache, int):
+                if element == cache:
+                    continue
 
-        if isinstance(cache, tuple):
-            if element == cache[1]:
-                continue
+                if element - cache == 1:
+                    cache = (cache, element)
+                    length += 1
+                    continue
 
-            if element - cache[1] == 1:
-                cache = (cache[0], element)
-                continue
+            if isinstance(cache, tuple):
+                if element == cache[1]:
+                    continue
+
+                if element - cache[1] == 1:
+                    cache = (cache[0], element)
+                    length += 1
+                    continue
 
-        yield cache
-        cache = element
+            yield cache
+            cache = element
+            length += 1
 
-    if cache is not None:
-        yield cache
+        if cache is not None:
+            yield cache
+
+        self.__length = length
+
+    @property
+    def length(self) -> int:
+        if self.__length is None:
+            raise TypeError("Length is not available until the iterator is exhausted")
+        return self.__length
+
+
+def compress_ranges(iterable: "Iterable[int]") -> "Iterable[int | tuple[int, int]]":
+    return CompressRanges(iterable)
 
 
 def decompress_ranges(
@@ -125,7 +148,9 @@ def tokenize_backwards(stream: bytes) -> "Iterable[bytes]":
 
 class UidSet(Iterable[int]):
     def __init__(self, iterable: "Iterable[int]", *, compress=True):
-        _data = b"".join(encode_compressed_ranges(compress_ranges(iterable)))
+        compress_ranges_iterator = CompressRanges(iterable)
+        _data = b"".join(encode_compressed_ranges(compress_ranges_iterator))
+        self.__length = compress_ranges_iterator.length
 
         # TODO: use a generator here
 
@@ -161,6 +186,9 @@ def __reversed__(self) -> "Iterator[int]":
             )
         )
 
+    def __len__(self) -> int:
+        return self.__length
+
 
 def make_data(length: int, ratio: float) -> list[int]:
     return [i for i in range(1, int(length * 1 / ratio) + 1) if random.random() < ratio]
@@ -176,6 +204,7 @@ def main():
         start = time.monotonic()
         ten_uid_set = UidSet(ten_list, compress=False)
         end = time.monotonic()
+        assert len(ten_uid_set) == len(ten_list)
         print(f"Time to create: {end - start:.2f}")
         print("Uid set length:", asizeof(ten_uid_set))
         print(f"Proportion: {asizeof(ten_uid_set) / asizeof(ten_list):.2f}")

tests/test_uid_set.py

@@ -1,8 +1,8 @@
 import pytest
 
 from inbox.uid_set import (
+    CompressRanges,
     UidSet,
-    compress_ranges,
     decode_compressed_ranges,
     decompress_ranges,
     encode_compressed_ranges,
@@ -25,7 +25,9 @@
     ),
 )
 def test_compress_decompress_ranges(example, expected):
-    assert list(compress_ranges(example)) == expected
+    compress_ranges = CompressRanges(example)
+    assert list(compress_ranges) == expected
+    assert compress_ranges.length == len(example)
     assert list(decompress_ranges(expected)) == sorted(example)
 
 
@@ -54,6 +56,8 @@ def test_encode_compressed_ranges(compressed_ranges, encoded_ranges):
         [1, 2],
         [1, 2, 3],
         [3, 1, 2],
+        [1, 1, 1],
+        [1, 2, 2, 5],
         [1, 2, 4],
         [1, 3, 4],
         [1, 2, 4, 5],
@@ -62,4 +66,6 @@ def test_encode_compressed_ranges(compressed_ranges, encoded_ranges):
     ),
 )
 def test_uid_set(iterable):
-    assert set(UidSet(iterable)) == set(iterable)
+    uid_set = UidSet(iterable)
+    assert set(uid_set) == set(iterable)
+    assert len(uid_set) == len(set(iterable))