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