Support generators for the excluded values

Author: jcjones

.gitignore

@@ -1,5 +1,6 @@
 # Compiled python modules.
 *.pyc
+.venv
 
 # Setuptools distribution folder.
 /dist/
@@ -8,4 +9,4 @@
 /build/
 
 # Python egg metadata, regenerated from source files by setuptools.
-/*.egg-info
+/*.egg-info

filtercascade/__init__.py

@@ -115,13 +115,28 @@ def __init__(self, filters, error_rates=[0.02, 0.5], growth_factor=1.1,
         self.version = version
 
     def initialize(self, *, include, exclude):
-        log.debug("{} include and {} exclude".format(
-            len(include), len(exclude)))
+        """
+            Arg "exclude" is potentially larger than main memory, so it should
+            be assumed to be passed as a lazy-loading iterator. If it isn't,
+            that's fine. The "include" arg must fit in memory and should be
+            assumed to be a set.
+        """
+        try:
+            iter(exclude)
+        except TypeError as te:
+            raise TypeError("exclude is not iterable", te)
+        try:
+            len(include)
+        except TypeError as te:
+            raise TypeError("include is not a list", te)
+
+        include_len = len(include)
+
         depth = 1
         maxSequentialGrowthLayers = 3
         sequentialGrowthLayers = 0
 
-        while len(include) > 0:
+        while include_len > 0:
             starttime = datetime.datetime.utcnow()
             er = self.error_rates[-1]
             if depth < len(self.error_rates):
@@ -133,24 +148,23 @@ def initialize(self, *, include, exclude):
                     # min_filter_length large. This is important for the deep layers near the end.
                     Bloomer.filter_with_characteristics(
                         max(
-                            int(len(include) * self.growth_factor),
+                            int(include_len * self.growth_factor),
                             self.min_filter_length), er, depth))
             else:
                 # Filter already created for this layer. Check size and resize if needed.
                 required_size = Bloomer.calc_size(
-                    self.filters[depth - 1].nHashFuncs, len(include), er)
+                    self.filters[depth - 1].nHashFuncs, include_len, er)
                 if self.filters[depth - 1].size < required_size:
                     # Resize filter
                     self.filters[depth -
                                  1] = Bloomer.filter_with_characteristics(
-                                     int(len(include) * self.growth_factor),
+                                     int(include_len * self.growth_factor),
                                      er, depth)
                     log.info("Resized filter at {}-depth layer".format(depth))
             filter = self.filters[depth - 1]
             log.debug(
-                "Initializing the {}-depth layer. err={} include={} exclude={} size={} hashes={}"
-                .format(depth, er, len(include), len(exclude), filter.size,
-                        filter.nHashFuncs))
+                "Initializing the {}-depth layer. err={} include_len={} size={} hashes={}"
+                .format(depth, er, include_len, filter.size, filter.nHashFuncs))
             # loop over the elements that *should* be there. Add them to the filter.
             for elem in include:
                 filter.add(elem)
@@ -188,7 +202,8 @@ def initialize(self, *, include, exclude):
                     sequentialGrowthLayers = 0
 
             include, exclude = false_positives, include
-            if len(include) > 0:
+            include_len = len(include)
+            if include_len > 0:
                 depth = depth + 1
         # Filter characteristics loaded from meta file may result in unused layers.
         # Remove them.

filtercascade/test.py

@@ -1,29 +1,48 @@
-import unittest
 import filtercascade
+import hashlib
+import unittest
+from itertools import islice
+
 
 class MockFile(object):
     def __init__(self):
         self.data = b""
+
     def __len__(self):
         return len(self.data)
+
     def __getitem__(self, idx):
         return self.data[idx]
 
     def write(self, s):
         self.data = self.data + s
+
     def read(self):
         return self.data
+
     def flush(self):
         pass
 
+
 class SimpleToByteClass(object):
     def __init__(self, ordinal):
         self.o = ordinal
         self.method_called = False
+
     def to_bytes(self):
         self.method_called = True
         return self.o.to_bytes(1, "little")
 
+
+def predictable_serial_gen(end):
+    counter = 0
+    while counter < end:
+        counter += 1
+        m = hashlib.sha256()
+        m.update(counter.to_bytes(4, byteorder="big"))
+        yield m.hexdigest()
+
+
 class TestFilterCascade(unittest.TestCase):
     def assertBloomerEqual(self, b1, b2):
         self.assertEqual(b1.nHashFuncs, b2.nHashFuncs)
@@ -76,6 +95,33 @@ def test_fc_input_formats(self):
         self.assertFilterCascadeEqual(f1, f2)
         self.assertFilterCascadeEqual(f1, f3)
 
+    def test_fc_include_not_list(self):
+        f = filtercascade.FilterCascade([])
+        with self.assertRaises(TypeError):
+            f.initialize(include=predictable_serial_gen(1),
+                         exclude=predictable_serial_gen(1))
+
+    def test_fc_exclude_must_be_iterable(self):
+        f = filtercascade.FilterCascade([])
+        with self.assertRaises(TypeError):
+            f.initialize(include=[], exclude=list(1))
+
+    def test_fc_iterable(self):
+        f = filtercascade.FilterCascade([])
+
+        serials = predictable_serial_gen(500_000)
+        # revocations must be disjoint from the main set, so
+        # slice off a set and re-use the remainder
+        revocations = set(islice(serials, 3_000))
+
+        f.initialize(include=revocations,
+                     exclude=serials)
+
+        self.assertEqual(len(f.filters), 3)
+        self.assertEqual(f.filters[0].size, 81272)
+        self.assertEqual(f.filters[1].size, 14400)
+        self.assertEqual(f.filters[2].size, 14400)
+
 
 if __name__ == '__main__':
-    unittest.main()
+    unittest.main()