Rev 3957: Add report_activity to osutils.pumpfile in lp:///~jameinel/bzr/1.12-network-io

[John Arbash Meinel]

At lp:///~jameinel/bzr/1.12-network-io

------------------------------------------------------------
revno: 3957
revision-id: john at arbash-meinel.com-20090123212055-qsa1szqtrrxtwgdl
parent: pqm at pqm.ubuntu.com-20090123181416-tku4gdtorboy6d0y
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: 1.12-network-io
timestamp: Fri 2009-01-23 15:20:55 -0600
message:
  Add report_activity to osutils.pumpfile
-------------- next part --------------
=== modified file 'bzrlib/osutils.py'
--- a/bzrlib/osutils.py	2009-01-13 03:06:36 +0000
+++ b/bzrlib/osutils.py	2009-01-23 21:20:55 +0000
@@ -508,7 +508,8 @@
     return False
 
 
-def pumpfile(from_file, to_file, read_length=-1, buff_size=32768):
+def pumpfile(from_file, to_file, read_length=-1, buff_size=32768,
+             report_activity=None, direction='read'):
     """Copy contents of one file to another.
 
     The read_length can either be -1 to read to end-of-file (EOF) or
@@ -517,6 +518,10 @@
     The buff_size represents the maximum size for each read operation
     performed on from_file.
 
+    :param report_activity: Call this as bytes are read, see
+        Transport._report_activity
+    :param direction: Will be passed to report_activity
+
     :return: The number of bytes copied.
     """
     length = 0
@@ -530,6 +535,8 @@
             if not block:
                 # EOF reached
                 break
+            if report_activity is not None:
+                report_activity(len(block), direction)
             to_file.write(block)
 
             actual_bytes_read = len(block)
@@ -542,6 +549,8 @@
             if not block:
                 # EOF reached
                 break
+            if report_activity is not None:
+                report_activity(len(block), direction)
             to_file.write(block)
             length += len(block)
     return length

=== modified file 'bzrlib/tests/test_osutils.py'
--- a/bzrlib/tests/test_osutils.py	2009-01-13 03:06:36 +0000
+++ b/bzrlib/tests/test_osutils.py	2009-01-23 21:20:55 +0000
@@ -486,6 +486,34 @@
             message = "Data not equal.  Expected %d bytes, received %d."
             self.fail(message % (len(response_data), self.test_data_len))
 
+    def test_report_activity(self):
+        activity = []
+        def log_activity(length, direction):
+            activity.append((length, direction))
+        from_file = StringIO(self.test_data)
+        to_file = StringIO()
+        pumpfile(from_file, to_file, buff_size=500,
+                 report_activity=log_activity, direction='read')
+        self.assertEqual([(500, 'read'), (500, 'read'), (500, 'read'),
+                          (36, 'read')], activity)
+
+        from_file = StringIO(self.test_data)
+        to_file = StringIO()
+        del activity[:]
+        pumpfile(from_file, to_file, buff_size=500,
+                 report_activity=log_activity, direction='write')
+        self.assertEqual([(500, 'write'), (500, 'write'), (500, 'write'),
+                          (36, 'write')], activity)
+
+        # And with a limited amount of data
+        from_file = StringIO(self.test_data)
+        to_file = StringIO()
+        del activity[:]
+        pumpfile(from_file, to_file, buff_size=500, read_length=1028,
+                 report_activity=log_activity, direction='read')
+        self.assertEqual([(500, 'read'), (500, 'read'), (28, 'read')], activity)
+
+
 
 class TestPumpStringFile(TestCase):