Merge pull request #125 from facetoe/cpu_usage_graph

Cpu usage graph

docs/i3pystatus.core.rst

@@ -74,3 +74,11 @@ core Package
     :undoc-members:
     :show-inheritance:
 
+:mod:`color` Module
+------------------
+
+.. automodule:: i3pystatus.core.color
+    :members:
+    :undoc-members:
+    :show-inheritance:
+

i3pystatus/core/color.py

@@ -0,0 +1,61 @@
+from colour import Color
+
+
+class ColorRangeModule(object):
+    """
+    Class to dynamically generate and select colors.
+
+    Requires the PyPI package `colour`
+    """
+
+    start_color = "#00FF00"
+    end_color = 'red'
+
+    @staticmethod
+    def get_hex_color_range(start_color, end_color, quantity):
+        """
+        Generates a list of quantity Hex colors from start_color to end_color.
+
+        :param start_color: Hex or plain English color for start of range
+        :param end_color: Hex or plain English color for end of range
+        :param quantity: Number of colours to return
+        :return: A list of Hex color values
+        """
+        raw_colors = [c.hex for c in list(Color(start_color).range_to(Color(end_color), quantity))]
+        colors = []
+        for color in raw_colors:
+
+            # i3bar expects the full Hex value but for some colors the colour
+            # module only returns partial values. So we need to convert these colors to the full
+            # Hex value.
+            if len(color) == 4:
+                fixed_color = "#"
+                for c in color[1:]:
+                    fixed_color += c * 2
+                colors.append(fixed_color)
+            else:
+                colors.append(color)
+        return colors
+
+    def get_gradient(self, value, colors, upper_limit=100):
+        """
+        Map a value to a color
+        :param value: Some value
+        :return: A Hex color code
+        """
+        index = int(self.percentage(value, upper_limit))
+        if index >= len(colors):
+            return colors[-1]
+        elif index < 0:
+            return colors[0]
+        else:
+            return colors[index]
+
+    @staticmethod
+    def percentage(part, whole):
+        """
+        Calculate percentage
+        """
+        if whole == 0:
+            return 0
+        return 100 * float(part) / float(whole)

i3pystatus/cpu_usage_graph.py

@@ -0,0 +1,54 @@
+from i3pystatus.core.color import ColorRangeModule
+from i3pystatus.cpu_usage import CpuUsage
+from i3pystatus.core.util import make_graph
+
+
+class CpuUsageGraph(CpuUsage, ColorRangeModule):
+    """
+     Shows CPU usage as a Unicode graph.
+     The first output will be inacurate.
+
+     Depends on the PyPI colour module - https://pypi.python.org/pypi/colour/0.0.5
+
+     Linux only
+
+     Available formatters:
+
+     * {cpu_graph}   graph of cpu usage.
+     * {usage}       usage average of all cores
+     * {usage_cpu*}  usage of one specific core. replace "*" by core number starting at 0
+     * {usage_all}   usage of all cores separate. usess natsort when available(relevant for more than 10 cores)
+     """
+
+    settings = (
+        ("cpu", "cpu to monitor, choices are 'usage_cpu' for all or 'usage_cpu*'. R"
+                "eplace '*' by core number starting at 0."),
+        ("start_color", "Hex or English name for start of color range, eg '#00FF00' or 'green'"),
+        ("end_color", "Hex or English name for end of color range, eg '#FF0000' or 'red'")
+    )
+
+    graph_width = 15
+    format = '{cpu_graph}'
+    cpu = 'usage_cpu'
+
+    def init(self):
+        super().init()
+        self.cpu_readings = self.graph_width * [0]
+        self.colors = self.get_hex_color_range(self.start_color, self.end_color, int(100))
+
+    def run(self):
+        format_options = self.get_usage()
+        core_reading = format_options[self.cpu]
+
+        self.cpu_readings.insert(0, core_reading)
+        self.cpu_readings = self.cpu_readings[:self.graph_width]
+
+        graph = make_graph(self.cpu_readings, 100.0)
+        format_options.update({'cpu_graph': graph})
+
+        color = self.get_gradient(core_reading, self.colors)
+        self.output = {
+            "full_text": self.format.format_map(format_options),
+            'color': color
+        }
+

i3pystatus/network_graph.py

@@ -0,0 +1,69 @@
+# -*- coding: utf-8 -*-
+from i3pystatus.core.color import ColorRangeModule
+from i3pystatus.network_traffic import NetworkTraffic
+from i3pystatus.core.util import make_graph
+
+
+class NetworkGraph(NetworkTraffic, ColorRangeModule):
+    """
+    Shows Network activity as a Unicode graph
+
+    Linux only
+
+    Requires the PyPI packages `psutil` and `colour`.
+
+    Available formatters:
+
+    {kbs}               Float representing kb\s
+    {network_graph}     Unicode network graph
+
+    """
+    settings = (
+        ("format", "format string"),
+        ("graph_width", "Width of the graph"),
+        ("upper_limit", "Expected max kb/s. This value controls how the graph is drawn and in what color"),
+        ("graph_type", "Whether to draw the graph for input or output. "
+                       "Allowed values 'input' or 'output'"),
+        ("divisor", "divide all byte values by this value"),
+        ("interface", "Interface to watch, eg 'eth0'"),
+        ("start_color", "Hex or English name for start of color range, eg '#00FF00' or 'green'"),
+        ("end_color", "Hex or English name for end of color range, eg '#FF0000' or 'red'")
+    )
+
+    format = "{network_graph}{kbs}KB/s"
+    graph_type = 'input'
+
+    interval = 1
+    graph_width = 15
+    upper_limit = 150.0
+
+    def init(self):
+        self.colors = self.get_hex_color_range(self.start_color, self.end_color, int(self.upper_limit))
+        self.kbs_arr = [0.0] * self.graph_width
+
+    def run(self):
+        self.update_counters()
+        if not self.pnic_before:
+            return
+
+        if self.graph_type == 'input':
+            kbs = self.get_bytes_received()
+        elif self.graph_type == 'output':
+            kbs = self.get_bytes_sent()
+        else:
+            raise Exception("graph_type must be either 'input' or 'output'!")
+
+        # Cycle array by inserting at the start and chopping off the last element
+        self.kbs_arr.insert(0, kbs)
+        self.kbs_arr = self.kbs_arr[:self.graph_width]
+
+        color = self.get_gradient(kbs, self.colors, self.upper_limit)
+        network_graph = make_graph(self.kbs_arr, self.upper_limit)
+
+        self.output = {
+            "full_text": self.format.format(
+                network_graph=network_graph,
+                kbs="{0:.1f}".format(round(kbs, 2)).rjust(6)
+            ),
+            'color': color,
+        }

i3pystatus/network_traffic.py

@@ -31,15 +31,32 @@ class NetworkTraffic(IntervalModule):
     round_size = None
  
     pnic = None
-    def run(self):
-        pnic_before = self.pnic
+    pnic_before = None
+
+    def update_counters(self):
+        self.pnic_before = self.pnic
         self.pnic = psutil.net_io_counters(pernic=True)[self.interface]
-        if not pnic_before: return
+
+    def get_bytes_sent(self):
+        return (self.pnic.bytes_sent - self.pnic_before.bytes_sent) / self.divisor
+
+    def get_bytes_received(self):
+        return (self.pnic.bytes_recv - self.pnic_before.bytes_recv) / self.divisor
+
+    def get_packets_sent(self):
+        return self.pnic.packets_sent - self.pnic_before.packets_sent
+
+    def get_packets_received(self):
+        return self.pnic.packets_recv - self.pnic_before.packets_recv
+
+    def run(self):
+        self.update_counters()
+        if not self.pnic_before: return
         cdict = {
-            "bytes_sent": (self.pnic.bytes_sent - pnic_before.bytes_sent) / self.divisor,
-            "bytes_recv": (self.pnic.bytes_recv - pnic_before.bytes_recv) / self.divisor,
-            "packets_sent": self.pnic.packets_sent - pnic_before.packets_sent,
-            "packets_recv": self.pnic.packets_recv - pnic_before.packets_recv,
+            "bytes_sent": self.get_bytes_sent(),
+            "bytes_recv": self.get_bytes_received(),
+            "packets_sent": self.get_packets_sent(),
+            "packets_recv": self.get_packets_received(),
         }
         round_dict(cdict, self.round_size)
         cdict["interface"] = self.interface