Added TopSamples utility to store the top N objects with the highest priorities.

Author: Mimino666

crawlmi/stats/top_samples.py

@@ -0,0 +1,32 @@
+from crawlmi.queue import Heap
+
+
+class TopSamples(object):
+    '''Stores the top `size` objects with the highest priorities.
+    Useful to store the top N largest sites, for example.
+    '''
+
+    def __init__(self, size=5):
+        self.size = size
+        self._heap = Heap()
+
+    def add_sample(self, priority, value):
+        # negate priority, so that the smallest ones are on top and first to pop
+        self._heap.push((-priority, value))
+        if len(self._heap) > self.size:
+            self._heap.pop()
+
+    @property
+    def samples(self):
+        objects = []
+        while self._heap:
+            objects.append(self._heap.pop())
+        map(self._heap.push, objects)
+        return [(-p, v) for (p, v) in reversed(objects)]
+
+    def __len__(self):
+        return len(self._heap)
+
+    def __str__(self):
+        return '\n    '.join(map(str, self.samples))
+    __repr__ = __str__

crawlmi/tests/test_top_samples.py

@@ -0,0 +1,31 @@
+import random
+
+from twisted.trial import unittest
+
+from crawlmi.stats.top_samples import TopSamples
+
+
+class TopSamplesTest(unittest.TestCase):
+    def test_size_1(self):
+        ts = TopSamples(size=1)
+        ts.add_sample(1, 'hello')
+        self.assertEqual(len(ts), 1)
+        self.assertListEqual(ts.samples, [(1, 'hello')])
+        ts.add_sample(0, 'world')
+        self.assertEqual(len(ts), 1)
+        self.assertListEqual(ts.samples, [(1, 'hello')])
+        ts.add_sample(2, '!')
+        self.assertEqual(len(ts), 1)
+        self.assertListEqual(ts.samples, [(2, '!')])
+
+    def test_default_size(self):
+        ts = TopSamples()
+        objects = [(i, 2 * i) for i in xrange(20)]
+        random.shuffle(objects)
+        for (p, v) in objects:
+            ts.add_sample(p, v)
+        self.assertEqual(len(ts), 5)
+        self.assertListEqual(ts.samples, [(19, 38), (18, 36), (17, 34), (16, 32), (15, 30)])
+        # test - requesting the values doesn't screw anything
+        self.assertEqual(len(ts), 5)
+        self.assertListEqual(ts.samples, [(19, 38), (18, 36), (17, 34), (16, 32), (15, 30)])