From f52702fbb59ce5b660e3588433744fd016b6d36a Mon Sep 17 00:00:00 2001
From: Tobias Witek <tobi@tobi-wan-kenobi.at>
Date: Fri, 5 May 2017 20:58:15 +0200
Subject: [PATCH] [modules/traffic] Refactor code for simplicity

Try to reduce code duplication by having a helper function that creates
the traffic widgets and do widget generation more generically.
---
 bumblebee/modules/traffic.py | 55 ++++++++++++++++--------------------
 1 file changed, 25 insertions(+), 30 deletions(-)

diff --git a/bumblebee/modules/traffic.py b/bumblebee/modules/traffic.py
index 73ef673..0da9fd1 100644
--- a/bumblebee/modules/traffic.py
+++ b/bumblebee/modules/traffic.py
@@ -33,45 +33,40 @@ class Module(bumblebee.engine.Module):
     def update(self, widgets):
         self._update_widgets(widgets)
 
+    def create_widget(self, widgets, name, txt=None, attributes={}):
+        widget = self.widget(name)
+        if widget: return widget
+
+        widget = bumblebee.output.Widget(name=name)
+        widget.full_text(txt)
+        widgets.append(widget)
+
+        for key in attributes:
+            widget.set(key, attributes[key])
+
+        return widget
+
     def _update_widgets(self, widgets):
         interfaces = [ i for i in netifaces.interfaces() if not i.startswith(self._exclude) ]
 
         counters = psutil.net_io_counters(pernic=True)
         for interface in interfaces:
             if not interface: interface = "lo"
-            rx = counters[interface].bytes_recv
-            tx = counters[interface].bytes_sent
+            data = {
+                "rx": counters[interface].bytes_recv,
+                "tx": counters[interface].bytes_sent,
+            }
 
             name = "traffic-{}".format(interface)
-            txname = "traffic.tx-{}".format(interface)
-            rxname = "traffic.rx-{}".format(interface)
 
-            widget = self.widget(name)
-            if not widget:
-                widget = bumblebee.output.Widget(name=name)
-                widgets.append(widget)
-                widget.full_text(interface)
+            self.create_widget(widgets, name, interface)
 
-            widget_rx = self.widget(rxname)
-            widget_tx = self.widget(txname)
-            if not widget_rx:
-                widget_rx = bumblebee.output.Widget(name=rxname)
-                widget_rx.set("theme.minwidth", "1000.00MB")
-                widgets.append(widget_rx)
-            if not widget_tx:
-                widget_tx = bumblebee.output.Widget(name=txname)
-                widget_tx.set("theme.minwidth", "1000.00MB")
-                widgets.append(widget_tx)
-
-            prev_rx = widget_rx.get("rx", 0)
-            prev_tx = widget_tx.get("tx", 0)
-            rxspeed = bumblebee.util.bytefmt(int(rx) - int(prev_rx))
-            txspeed = bumblebee.util.bytefmt(int(tx) - int(prev_tx))
-
-            widget_rx.full_text(rxspeed)
-            widget_tx.full_text(txspeed)
-
-            widget_rx.set("rx", rx)
-            widget_tx.set("tx", tx)
+            for direction in ["rx", "tx"]:
+                name = "traffic.{}-{}".format(direction, interface)
+                widget = self.create_widget(widgets, name, attributes={"theme.minwidth": "1000.00MB"})
+                prev = widget.get(direction, 0)
+                speed = bumblebee.util.bytefmt(int(data[direction]) - int(prev))
+                widget.full_text(speed)
+                widget.set(direction, data[direction])
 
 # vim: tabstop=8 expandtab shiftwidth=4 softtabstop=4