Allow ByteSpeed to show current-ish speed

3 files changed, 52 insertions, 12 deletions

diff --git a/NEWS b/NEWS
index c6fe4d7..9af320f 100644
--- a/NEWS
+++ b/NEWS
@@ -1,6 +1,13 @@
 NEWS file for ttystatus
 =======================
 
+Version 0.18, released UNRELEASED
+---------------------------------
+
+* `ByteSpeed` can now show average speed over a short period of time
+  (which approximates current speed), instead of over the whole duration 
+  of the program.
+
 Version 0.17, released 2012-04-15
 ---------------------------------
 
diff --git a/ttystatus/bytespeed.py b/ttystatus/bytespeed.py
index 3331069..1e619f2 100644
--- a/ttystatus/bytespeed.py
+++ b/ttystatus/bytespeed.py
@@ -1,4 +1,4 @@
-# Copyright 2010  Lars Wirzenius
+# Copyright 2010, 2012  Lars Wirzenius
 # 
 # This program is free software: you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -25,14 +25,13 @@ class ByteSpeed(ttystatus.Widget):
 
     static_width = False
     
-    def __init__(self, name):
+    def __init__(self, name, duration=None):
         self.name = name
-        self._bytes = 0
-        self._started = None
+        self._duration = None if duration is None else float(duration)
+        self._data_points = []
         
-    def now(self):
+    def now(self): # pragma: no cover
         '''Wrapper around time.time for unit tests to overrride.'''
-        
         return time.time()
         
     def render(self, width):
@@ -43,11 +42,14 @@ class ByteSpeed(ttystatus.Widget):
             (1024**1, 1, 'KiB/s'),
         )
         
-        if self._started is None:
+        if len(self._data_points) < 2:
             return '0 B/s'
         
-        duration = self.now() - self._started
-        speed = self._bytes / duration
+        oldest_bytes, started = self._data_points[0]
+        latest_bytes, dummy = self._data_points[-1]
+        bytes = latest_bytes - oldest_bytes
+        duration = self.now() - started
+        speed = bytes / duration
         
         for factor, decimals, unit in units:
             if speed >= factor:
@@ -57,6 +59,14 @@ class ByteSpeed(ttystatus.Widget):
         return '%.0f B/s' % speed
         
     def update(self, master):
-        if self._started is None:
-            self._started = self.now()
-        self._bytes = master[self.name]
+        bytes = master[self.name]
+        now = self.now()
+        self._data_points.append((bytes, now))
+        if self._duration is None:
+            if len(self._data_points) > 2:
+                del self._data_points[1:-1]
+        else:
+            cutoff = now - self._duration
+            while self._data_points[0][1] < cutoff:
+                del self._data_points[0]
+
diff --git a/ttystatus/bytespeed_tests.py b/ttystatus/bytespeed_tests.py
index 9bcc1e0..437e32b 100644
--- a/ttystatus/bytespeed_tests.py
+++ b/ttystatus/bytespeed_tests.py
@@ -31,6 +31,9 @@ class ByteSpeedTests(unittest.TestCase):
         self.assertEqual(self.w.render(0), '0 B/s')
 
     def test_formats_zero_bytes_correctly(self):
+        self.w.now = lambda: 1
+        self.w.update({ 'foo': 0 })
+        self.w.now = lambda: 2
         self.w.update({ 'foo': 0 })
         self.assertEqual(self.w.render(0), '0 B/s')
 
@@ -55,3 +58,23 @@ class ByteSpeedTests(unittest.TestCase):
         self.w.update({ 'foo': 10 * 1024**4 })
         self.assertEqual(self.w.render(0), '10.00 TiB/s')
 
+    def test_keeps_only_two_data_points_with_infinite_duration(self):
+        for when in range(100):
+            self.w.now = lambda: when
+            self.w.update({ 'foo': 0 })
+        self.assertEqual(self.w.render(0), '0 B/s')
+
+    def test_shows_current_speed_when_requested(self):
+        items = [
+            (0, 0),
+            (1, 1024),
+            (10, 1024),
+            (11, 1024),
+        ]
+
+        w = ttystatus.ByteSpeed('foo', duration=5)
+        for when, bytes in items:
+            w.now = lambda: when
+            w.update({ 'foo': bytes })
+        self.assertEqual(w.render(0), '0 B/s')
+