[modules/traffic] compute theme.minwidth based on traffic.format

bumblebee/modules/traffic.py

@@ -10,6 +10,7 @@ Parameters:
                       Defaults to "{:.2f}"
 """
 
+import re
 import time
 import psutil
 import netifaces
@@ -70,6 +71,20 @@ class Module(bumblebee.engine.Module):
             return []
         return retval
 
+    def get_minwidth_str(self):
+        """computes theme.minwidth string based on traffic.format parameter"""
+        minwidth_str = "1000"
+        try:
+            length = int(re.match("{:\.(\d+)f}", self._format).group(1))
+            if length > 0:
+                minwidth_str += "." + "0" * length
+        except AttributeError:
+            # return default value
+            return "1000.00MB"
+        finally:
+            minwidth_str += "MB"
+        return minwidth_str
+
     def _update_widgets(self, widgets):
         interfaces = [i for i in netifaces.interfaces() if not i.startswith(self._exclude)]
 
@@ -103,7 +118,7 @@ class Module(bumblebee.engine.Module):
 
             for direction in ["rx", "tx"]:
                 name = "traffic.{}-{}".format(direction, interface)
-                widget = self.create_widget(widgets, name, attributes={"theme.minwidth": "1000.00MB"})
+                widget = self.create_widget(widgets, name, attributes={"theme.minwidth": self.get_minwidth_str()})
                 prev = self._prev.get(name, 0)
                 speed = bumblebee.util.bytefmt(
                     (int(data[direction]) - int(prev))/timediff,

tests/modules/test_traffic.py

@@ -0,0 +1,23 @@
+import mock
+import unittest
+
+import tests.mocks as mocks
+
+from bumblebee.modules.traffic import Module
+
+class TestTrafficModule(unittest.TestCase):
+    def setUp(self):
+        mocks.setup_test(self, Module)
+
+    def test_default_format(self):
+        self.assertEqual(self.module._format, "{:.2f}")
+
+    def test_get_minwidth_str(self):
+        # default value (two digits after dot)
+        self.assertEqual(self.module.get_minwidth_str(), "1000.00MB")
+        # integer value
+        self.module._format = "{:.0f}"
+        self.assertEqual(self.module.get_minwidth_str(), "1000MB")
+        # just one digit after dot
+        self.module._format = "{:.1f}"
+        self.assertEqual(self.module.get_minwidth_str(), "1000.0MB")