From 21e2f46564648cb473ec7d23a7b57fce03f7e00b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Thayn=C3=A3=20B=2E=20Moretti?= <moretty@gmail.com>
Date: Tue, 9 Jul 2019 22:27:59 -0300
Subject: [PATCH] Add network traffic icons & minor refactor

---
 bumblebee/modules/network_traffic.py | 60 +++++++++++++++++++++-------
 themes/icons/ascii.json              |  4 ++
 themes/icons/awesome-fonts.json      |  4 ++
 themes/icons/ionicons.json           |  4 ++
 4 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/bumblebee/modules/network_traffic.py b/bumblebee/modules/network_traffic.py
index ad819ec..a8aead2 100644
--- a/bumblebee/modules/network_traffic.py
+++ b/bumblebee/modules/network_traffic.py
@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 
 """Displays network traffic
+No extra configuratio needed.
 """
 
 import psutil
@@ -12,7 +13,11 @@ import bumblebee.output
 import bumblebee.engine
 import bumblebee.util
 
+WIDGET_NAME = 'network_traffic'
+
 class Module(bumblebee.engine.Module):
+    """Bumblebee main module """
+
     def __init__(self, engine, config):
         super(Module, self).__init__(engine, config)
 
@@ -21,9 +26,21 @@ class Module(bumblebee.engine.Module):
 
             self._download_tx = self._bandwidth.bytes_recv()
             self._upload_tx = self._bandwidth.bytes_sent()
-        except:
+        except Exception:
+            """ We do not want do explode anything """
             pass
 
+    @classmethod
+    def state(cls, widget):
+        """Return the widget state"""
+
+        if widget.name == '{}.rx'.format(WIDGET_NAME):
+            return 'rx'
+        elif widget.name == '{}.tx'.format(WIDGET_NAME):
+            return 'tx'
+
+        return None
+
     def update(self, widgets):
         try:
             download_tx = self._bandwidth.bytes_recv()
@@ -35,50 +52,63 @@ class Module(bumblebee.engine.Module):
             self.update_widgets(widgets, download_rate, upload_rate)
 
             self._download_tx, self._upload_tx = download_tx, upload_tx
-        except:
+        except Exception:
+            """ We do not want do explode anything """
             pass
 
-    def update_widgets(self, widgets, download_rate, upload_rate):
+    @classmethod
+    def update_widgets(cls, widgets, download_rate, upload_rate):
+        """Update tx/rx widgets with new rates"""
         del widgets[:]
 
         widgets.extend((
-            TrafficWidget(text=download_rate, icon=''),
-            TrafficWidget(text=upload_rate, icon='')
+            TrafficWidget(text=download_rate, direction='tx'),
+            TrafficWidget(text=upload_rate, direction='rx')
         ))
 
 
 class BandwidthInfo(object):
+    """Get received/sent bytes from network adapter"""
+
     def bytes_recv(self):
+        """Return received bytes"""
         return self.bandwidth().bytes_recv
 
     def bytes_sent(self):
+        """Return sent bytes"""
         return self.bandwidth().bytes_sent
 
     def bandwidth(self):
+        """Return bandwidth information"""
         io_counters = self.io_counters()
         return io_counters[self.default_network_adapter()]
 
-    def default_network_adapter(self):
+    @classmethod
+    def default_network_adapter(cls):
+        """Return default active network adapter"""
         gateway = netifaces.gateways()['default']
 
-        if (len(gateway) == 0):
+        if not gateway:
             raise 'No default gateway found'
 
         return gateway[netifaces.AF_INET][1]
 
-    def io_counters(self):
+    @classmethod
+    def io_counters(cls):
+        """Return IO counters"""
         return psutil.net_io_counters(pernic=True)
 
-
 class TrafficWidget(object):
-    def __new__(self, text, icon):
-        widget = bumblebee.output.Widget()
-        widget.set('theme.minwidth', '00000000KiB/s')
-        widget.full_text(self.humanize(text, icon))
+    """Create a traffic widget with humanized bytes string with proper icon (up/down)"""
+    def __new__(cls, text, direction):
+        widget = bumblebee.output.Widget(name='{0}.{1}'.format(WIDGET_NAME, direction))
+        widget.set('theme.minwidth', '0000000KiB/s')
+        widget.full_text(cls.humanize(text))
 
         return widget
 
     @staticmethod
-    def humanize(text, icon):
+    def humanize(text):
+        """Return humanized bytes"""
         humanized_byte_format = bumblebee.util.bytefmt(text)
-        return '{0} {1}/s'.format(icon, humanized_byte_format)
+        return '{0}/s'.format(humanized_byte_format)
diff --git a/themes/icons/ascii.json b/themes/icons/ascii.json
index 6358442..9f345b5 100644
--- a/themes/icons/ascii.json
+++ b/themes/icons/ascii.json
@@ -92,6 +92,10 @@
 		"rx": { "prefix": "down"},
 		"tx": { "prefix": "up"}
 	},
+	"network_traffic": {
+		"rx": { "prefix": "down" },
+		"tx": { "prefix": "up" }
+	},
     "mpd": {
 		"playing": { "prefix": ">" },
 		"paused": { "prefix": "||" },
diff --git a/themes/icons/awesome-fonts.json b/themes/icons/awesome-fonts.json
index 5513466..f436179 100644
--- a/themes/icons/awesome-fonts.json
+++ b/themes/icons/awesome-fonts.json
@@ -145,6 +145,10 @@
 	    "rx": { "prefix": "" },
 	    "tx": { "prefix": "" }
     },
+    "network_traffic": {
+	    "rx": { "prefix": "" },
+            "tx": { "prefix": "" }
+    },
     "mpd": {
 		"playing": { "prefix": "" },
 		"paused": { "prefix": "" },
diff --git a/themes/icons/ionicons.json b/themes/icons/ionicons.json
index f2f8fc1..bd9ee4b 100644
--- a/themes/icons/ionicons.json
+++ b/themes/icons/ionicons.json
@@ -129,6 +129,10 @@
 	  "rx": { "prefix": "\uf365" },
 	  "tx": { "prefix": "\uf35f" }
   },
+  "network_traffic": {
+	  "rx": { "prefix": "\uf365" },
+	  "tx": { "prefix": "\uf35f" }
+  },
   "mpd": {
 		"playing": { "prefix": "\uf488" },
 		"paused": { "prefix": "\uf210" },