Merge pull request #110 from FrnkPsycho/traffic_widget

feat: add network traffic widget and some fixes

src/config.yaml

@@ -458,6 +458,14 @@ widgets:
           on_middle: "do_nothing"
           on_right: "exec cmd.exe /c start ms-settings:network"
 
+  traffic:
+  type: "yasb.traffic.TrafficWidget"
+  options:
+      label: "\ueb01 \ueab4 {download_speed} | \ueab7 {upload_speed}"
+      label_alt: "\ueb01 \ueab4 {upload_speed} | \ueab7 {download_speed}"
+      update_interval: 1000 # Update interval should be a multiple of 1000
+      callbacks:
+        on_right: "exec cmd /c Taskmgr"
 
   # Some custom widgets
 

src/core/validation/widgets/yasb/traffic.py

@@ -0,0 +1,39 @@
+DEFAULTS = {
+    "label": "\ueb01 \ueab4 {download_speed} | \ueab7 {upload_speed}",
+    "label_alt": "\ueb01 \ueab4 {upload_speed} | \ueab7 {download_speed}",
+    "update_interval": 1000,
+    "callbacks": {
+        "on_left": "toggle_label",
+        "on_middle": "do_nothing",
+        "on_right": "do_nothing",
+    },
+}
+
+VALIDATION_SCHEMA = {
+    "label": {"type": "string", "default": DEFAULTS["label"]},
+    "label_alt": {"type": "string", "default": DEFAULTS["label_alt"]},
+    "update_interval": {
+        "type": "integer",
+        "default": DEFAULTS["update_interval"],
+        "min": 0,
+        "max": 60000,
+    },
+    "callbacks": {
+        "type": "dict",
+        "schema": {
+            "on_left": {
+                "type": "string",
+                "default": DEFAULTS["callbacks"]["on_left"],
+            },
+            "on_middle": {
+                "type": "string",
+                "default": DEFAULTS["callbacks"]["on_middle"],
+            },
+            "on_right": {
+                "type": "string",
+                "default": DEFAULTS["callbacks"]["on_right"],
+            },
+        },
+        "default": DEFAULTS["callbacks"],
+    },
+}

src/core/validation/widgets/yasb/wifi.py

@@ -1,6 +1,4 @@
 DEFAULTS = {
-    # 'label': '\uf017 {%H:%M:%S}',
-    # 'label_alt': '\uf017 {%d-%m-%y %H:%M:%S}',
     'label': "{wifi_icon}",
     'label_alt': "{wifi_icon} {wifi_name}",
     'update_interval': 1000,

src/core/widgets/example.py

@@ -24,7 +24,9 @@ def __init__(
         self._label_alt.setProperty("class", "label alt")
         self.widget_layout.addWidget(self._label)
         self.widget_layout.addWidget(self._label_alt)
+
         self.register_callback("toggle_label", self._toggle_label)
+        self.register_callback("update_label", self._update_label)
 
         self.callback_left = callbacks['on_left']
         self.callback_right = callbacks['on_right']

src/core/widgets/yasb/traffic.py

@@ -0,0 +1,105 @@
+import psutil
+from humanize import naturalsize
+from core.widgets.base import BaseWidget
+from core.validation.widgets.yasb.traffic import VALIDATION_SCHEMA
+from PyQt6.QtWidgets import QLabel
+
+
+class TrafficWidget(BaseWidget):
+    validation_schema = VALIDATION_SCHEMA
+
+    # initialize io counters
+    io = psutil.net_io_counters()
+    bytes_sent = io.bytes_sent
+    bytes_recv = io.bytes_recv
+    interval = 1  # second(s)
+
+    def __init__(
+        self,
+        label: str,
+        label_alt: str,
+        update_interval: int,
+        callbacks: dict[str, str],
+    ):
+        super().__init__(update_interval, class_name="traffic-widget")
+        self.interval = update_interval // 1000
+
+        self._show_alt_label = False
+        self._label_content = label
+        self._label_alt_content = label_alt
+
+        self._label = QLabel()
+        self._label_alt = QLabel()
+        self._label.setProperty("class", "label")
+        self._label_alt.setProperty("class", "label alt")
+        self.widget_layout.addWidget(self._label)
+        self.widget_layout.addWidget(self._label_alt)
+
+        self.register_callback("toggle_label", self._toggle_label)
+        self.register_callback("update_label", self._update_label)
+
+        self.callback_left = callbacks["on_left"]
+        self.callback_right = callbacks["on_right"]
+        self.callback_middle = callbacks["on_middle"]
+        self.callback_timer = "update_label"
+
+        self._label.show()
+        self._label_alt.hide()
+
+        self.start_timer()
+
+    def _toggle_label(self):
+        self._show_alt_label = not self._show_alt_label
+
+        if self._show_alt_label:
+            self._label.hide()
+            self._label_alt.show()
+        else:
+            self._label.show()
+            self._label_alt.hide()
+
+        self._update_label()
+
+    def _update_label(self):
+        # Update the active label at each timer interval
+        active_label = self._label_alt if self._show_alt_label else self._label
+        active_label_content = self._label_alt_content if self._show_alt_label else self._label_content
+        active_label_formatted = active_label_content
+
+        try:
+            upload_speed, download_speed = self._get_speed()
+        except Exception:
+            upload_speed, download_speed = "N/A", "N/A"
+
+        label_options = [
+            ("{upload_speed}", upload_speed),
+            ("{download_speed}", download_speed),
+        ]
+
+        for option, value in label_options:
+            active_label_formatted = active_label_formatted.replace(option, str(value))
+
+        active_label.setText(active_label_formatted)
+
+    def _get_speed(self) -> [str, str]:
+        current_io = psutil.net_io_counters()
+        upload_diff = current_io.bytes_sent - self.bytes_sent
+        download_diff = current_io.bytes_recv - self.bytes_recv
+
+        if upload_diff < 1024:
+            upload_speed = f"{upload_diff} B/s"
+        else:
+            upload_speed = naturalsize(
+                (current_io.bytes_sent - self.bytes_sent) // self.interval,
+            ) + "/s"
+
+        if download_diff < 1024:
+            download_speed = f"{download_diff} B/s"
+        else:
+            download_speed = naturalsize(
+                (current_io.bytes_recv - self.bytes_recv) // self.interval,
+            ) + "/s"
+
+        self.bytes_sent = current_io.bytes_sent
+        self.bytes_recv = current_io.bytes_recv
+        return upload_speed, download_speed

src/core/widgets/yasb/wifi.py

@@ -8,12 +8,12 @@ class WifiWidget(BaseWidget):
     validation_schema = VALIDATION_SCHEMA
 
     def __init__(
-            self,
-            label: str,
-            label_alt: str,
-            update_interval: int,
-            wifi_icons: list[str],
-            callbacks: dict[str, str],
+        self,
+        label: str,
+        label_alt: str,
+        update_interval: int,
+        wifi_icons: list[str],
+        callbacks: dict[str, str],
     ):
         super().__init__(update_interval, class_name="wifi-widget")
         self._wifi_icons = wifi_icons
@@ -32,9 +32,9 @@ def __init__(
         self.register_callback("toggle_label", self._toggle_label)
         self.register_callback("update_label", self._update_label)
 
-        self.callback_left = callbacks['on_left']
-        self.callback_right = callbacks['on_right']
-        self.callback_middle = callbacks['on_middle']
+        self.callback_left = callbacks["on_left"]
+        self.callback_right = callbacks["on_right"]
+        self.callback_middle = callbacks["on_middle"]
         self.callback_timer = "update_label"
 
         self._label.show()
@@ -61,33 +61,40 @@ def _update_label(self):
         # Determine which label is active
         active_label = self._label_alt if self._show_alt_label else self._label
 
-        if self._show_alt_label:
-            updated_content = f"{wifi_icon} {wifi_name}"
-        else:
-            updated_content = f"{wifi_icon}"
+        label_options = [
+            ("{wifi_icon}", wifi_icon),
+            ("{wifi_name}", wifi_name),
+        ]
+
+        # Format the label content
+        updated_content = self._label_alt_content if self._show_alt_label else self._label_content
+        for label_option in label_options:
+            updated_content = updated_content.replace(
+                label_option[0], str(label_option[1])
+            )
 
         active_label.setText(updated_content)
 
     def _get_wifi_strength(self):
         # Get the wifi strength from the system
-        result = result = os.popen('netsh wlan show interfaces').read()
+        result = os.popen("netsh wlan show interfaces").read()
 
         # Return 0 if no wifi interface is found
         if "There is no wireless interface on the system." in result:
             return 0
 
         # Extract signal strength from the result
-        for line in result.split('\n'):
-            if "Signal" in line:
-                strength = line.split(":")[1].strip().split(' ')[0].replace('%', '')
+        for line in result.split("\n"):
+            if "Signal" in line:  # FIXME: This will break if the system language is not English
+                strength = line.split(":")[1].strip().split(" ")[0].replace("%", "")
                 return int(strength)
 
         return 0
 
     def _get_wifi_name(self):
-        result = result = os.popen('netsh wlan show interfaces').read()
+        result = os.popen("netsh wlan show interfaces").read()
 
-        for line in result.split('\n'):
+        for line in result.split("\n"):
             if "SSID" in line:
                 return line.split(":")[1].strip()
 

src/styles.css

@@ -75,12 +75,17 @@
 .cpu-widget .label,
 .wifi-widget .label,
 .memory-widget .label,
+.traffic-widget .label,
 .battery-widget .label {
   padding: 2px 5px;
   font-size: 14px;
   border-radius: 5px;
 }
 
+.traffic-widget .label {
+    background: #B3D3D9;
+}
+
 .wifi-widget .label {
   background: #b8bb26;
 }