A module for displaying NVIDIA GPUs' temperature

1 files changed, 80 insertions, 0 deletions

diff --git a/py3status/modules/nvidia_temp.py b/py3status/modules/nvidia_temp.py
new file mode 100755
index 0000000..4f6a40e
--- /dev/null
+++ b/py3status/modules/nvidia_temp.py
@@ -0,0 +1,80 @@
+#!/usr/bin/env python
+# encoding: utf-8
+
+"""
+Module for displaying NVIDIA GPU temperature
+
+Requires:
+    - nvidia-smi
+
+@author jmdana <https://github.com/jmdana>
+@license GPLv3 <http://www.gnu.org/licenses/gpl-3.0.txt>
+"""
+
+import re
+import shlex
+from subprocess import check_output
+from time import time
+
+TEMP_RE = re.compile(r"Current Temp\s+:\s+([0-9]+)")
+
+class Py3status:
+    # configuration parameters
+    cache_timeout = 10
+    color_good = None
+    color_bad = None
+    format_prefix = "GPU: "
+    format_units = "°C"
+    separator = '|'
+
+    def nvidia_temp(self, i3s_output_list, i3s_config):
+        # The whole command:
+        # nvidia-smi -q -d TEMPERATURE | sed -nr 's/.*Current Temp.*:[[:space:]]*([0-9]+).*/\1/p'
+
+        out = check_output(shlex.split("nvidia-smi -q -d TEMPERATURE"))
+        temps = re.findall(TEMP_RE, out.decode("utf-8"))
+
+        if temps != []:
+            data = []
+            for temp in temps:
+                fmt_str = "{temp}{format_units}".format(
+                    temp=temp,
+                    format_units=self.format_units
+                )
+                data.append(fmt_str)
+
+            output = "{format_prefix}{data}".format(
+                format_prefix=self.format_prefix,
+                data=self.separator.join(data)
+            )
+
+            color = self.color_good or i3s_config['color_good']
+        else:
+            output = "{format_prefix}OFF".format(
+                format_prefix=self.format_prefix
+            )
+
+            color = self.color_bad or i3s_config['color_bad']
+
+        response = {
+            'cached_until': time() + self.cache_timeout,
+            'full_text': output,
+            'color': color,
+        }
+
+        return response
+
+if __name__ == "__main__":
+    """
+    Test this module by calling it directly.
+    """
+    from time import sleep
+    x = Py3status()
+    config = {
+        "color_good": "#00FF00",
+        "color_bad": "#FF0000",
+        }
+
+    while True:
+        print(x.nvidia_temp([], config))
+        sleep(1)