Rev 3874: Add a VFDecorator that can yield records in a specified order in http://bzr.arbash-meinel.com/branches/bzr/1.10-dev/304841-not-present-chain

[John Arbash Meinel]

At http://bzr.arbash-meinel.com/branches/bzr/1.10-dev/304841-not-present-chain

------------------------------------------------------------
revno: 3874
revision-id: john at arbash-meinel.com-20081203210501-fejpuwpvfkvwogxd
parent: pqm at pqm.ubuntu.com-20081202062535-khqslu9vyrvzr7om
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: 304841-not-present-chain
timestamp: Wed 2008-12-03 15:05:01 -0600
message:
  Add a VFDecorator that can yield records in a specified order
-------------- next part --------------
=== modified file 'bzrlib/tests/test_versionedfile.py'
--- a/bzrlib/tests/test_versionedfile.py	2008-07-17 09:13:16 +0000
+++ b/bzrlib/tests/test_versionedfile.py	2008-12-03 21:05:01 +0000
@@ -2218,3 +2218,58 @@
         record = it.next()
         self.assertEquals("absent", record.storage_kind)
 
+
+class TestOrderingVersionedFilesDecorator(TestCaseWithMemoryTransport):
+
+    def get_ordering_vf(self, key_priority):
+        builder = self.make_branch_builder('test')
+        builder.start_series()
+        builder.build_snapshot('A', None, [
+            ('add', ('', 'TREE_ROOT', 'directory', None))])
+        builder.build_snapshot('B', ['A'], [])
+        builder.build_snapshot('C', ['B'], [])
+        builder.build_snapshot('D', ['C'], [])
+        builder.finish_series()
+        b = builder.get_branch()
+        b.lock_read()
+        self.addCleanup(b.unlock)
+        vf = b.repository.inventories
+        return versionedfile.OrderingVersionedFilesDecorator(vf, key_priority)
+
+    def test_get_empty(self):
+        vf = self.get_ordering_vf({})
+        self.assertEqual([], vf.calls)
+
+    def test_get_record_stream_topological(self):
+        vf = self.get_ordering_vf({('A',): 3, ('B',): 2, ('C',): 4, ('D',): 1})
+        request_keys = [('B',), ('C',), ('D',), ('A',)]
+        keys = [r.key for r in vf.get_record_stream(request_keys,
+                                    'topological', False)]
+        # We should have gotten the keys in topological order
+        self.assertEqual([('A',), ('B',), ('C',), ('D',)], keys)
+        # And recorded that the request was made
+        self.assertEqual([('get_record_stream', request_keys, 'topological',
+                           False)], vf.calls)
+
+    def test_get_record_stream_ordered(self):
+        vf = self.get_ordering_vf({('A',): 3, ('B',): 2, ('C',): 4, ('D',): 1})
+        request_keys = [('B',), ('C',), ('D',), ('A',)]
+        keys = [r.key for r in vf.get_record_stream(request_keys,
+                                   'unordered', False)]
+        # They should be returned based on their priority
+        self.assertEqual([('D',), ('B',), ('A',), ('C',)], keys)
+        # And the request recorded
+        self.assertEqual([('get_record_stream', request_keys, 'unordered',
+                           False)], vf.calls)
+
+    def test_get_record_stream_implicit_order(self):
+        vf = self.get_ordering_vf({('B',): 2, ('D',): 1})
+        request_keys = [('B',), ('C',), ('D',), ('A',)]
+        keys = [r.key for r in vf.get_record_stream(request_keys,
+                                   'unordered', False)]
+        # A and C are not in the map, so they get sorted to the front. A comes
+        # before C alphabetically, so it comes back first
+        self.assertEqual([('A',), ('C',), ('D',), ('B',)], keys)
+        # And the request recorded
+        self.assertEqual([('get_record_stream', request_keys, 'unordered',
+                           False)], vf.calls)

=== modified file 'bzrlib/versionedfile.py'
--- a/bzrlib/versionedfile.py	2008-11-21 20:20:52 +0000
+++ b/bzrlib/versionedfile.py	2008-12-03 21:05:01 +0000
@@ -521,7 +521,7 @@
     """
 
     def __init__(self, backing_vf):
-        """Create a RecordingVersionedFileDsecorator decorating backing_vf.
+        """Create a RecordingVersionedFilesDecorator decorating backing_vf.
         
         :param backing_vf: The versioned file to answer all methods.
         """
@@ -562,6 +562,44 @@
         return self._backing_vf.keys()
 
 
+class OrderingVersionedFilesDecorator(RecordingVersionedFilesDecorator):
+    """A VF that records calls, and returns keys in specific order.
+
+    :ivar calls: A list of the calls made; can be reset at any time by
+        assigning [] to it.
+    """
+
+    def __init__(self, backing_vf, key_priority):
+        """Create a RecordingVersionedFilesDecorator decorating backing_vf.
+
+        :param backing_vf: The versioned file to answer all methods.
+        :param key_priority: A dictionary defining what order keys should be
+            returned from an 'unordered' get_record_stream request.
+            Keys with lower priority are returned first, keys not present in
+            the map get an implicit priority of 0, and are returned in
+            lexicographical order.
+        """
+        RecordingVersionedFilesDecorator.__init__(self, backing_vf)
+        self._key_priority = key_priority
+
+    def get_record_stream(self, keys, sort_order, include_delta_closure):
+        self.calls.append(("get_record_stream", list(keys), sort_order,
+            include_delta_closure))
+        if sort_order == 'unordered':
+            def sort_key(key):
+                return (self._key_priority.get(key, 0), key)
+            # Use a defined order by asking for the keys one-by-one from the
+            # backing_vf
+            for key in sorted(keys, key=sort_key):
+                for record in self._backing_vf.get_record_stream([key],
+                                'unordered', include_delta_closure):
+                    yield record
+        else:
+            for record in self._backing_vf.get_record_stream(keys, sort_order,
+                            include_delta_closure):
+                yield record
+
+
 class KeyMapper(object):
     """KeyMappers map between keys and underlying partitioned storage."""